diff --git a/.gitignore b/.gitignore index 1978c22..3ffe08f 100644 --- a/.gitignore +++ b/.gitignore @@ -3,23 +3,40 @@ data/ job_scripts/ outputs/ profs/ +sparse-extension/build.out sparse-extension/build/ sparse-extension/dist/ sparse-extension/sparse_coo_tensor_cpp.egg-info/ sparse_coo_tensor_cpp.egg-info/ .*.swp +.*.swo +.*.swm +.*.swn .nfs* *.txt +!requirements.txt *.pt -*.qdrep -tests/ -job* slurm_outputs/ -examples/cagnet_outputs/ +job-* +job_* +gcn_distr.py +gcn_distr_15d.py +gcn_distr_2d.py +gcn_distr_3d.py +gcn.py +reddit.py +tests/data +tests/dgl/ +tests/examples +tests/LADIES/ +tests/test_nccl +tests/pytorch_geometric/ +tests/torch-quiver/ +tests/*.out +examples/wandb +examples/*.batch +examples/*.png +*.qdrep *.out -examples/*.nsys-rep -examples/*.ncu-rep -cagnet/ -examples/ -slurm_outputs/ -tests/ +*.nsys-rep +*.ncu-rep diff --git a/README.md b/README.md index 660bc5a..33f8685 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,5 @@ # CAGNET: Communication-Avoiding Graph Neural nETworks -This branch contains implementations for CAGNET's full-batch training pipeline (SC'20). For CAGNET's minibatch training pipeline (MLSys'24), please refer to the `distributed-sampling` branch. - ## Description CAGNET is a family of parallel algorithms for training GNNs that can asymptotically reduce communication compared to previous parallel GNN training methods. CAGNET algorithms are based on 1D, 1.5D, 2D, and 3D sparse-dense matrix multiplication, and are implemented with `torch.distributed` on GPU-equipped clusters. We also implement these parallel algorithms on a 2-layer GCN. @@ -12,13 +10,19 @@ For more information, please read our ACM/IEEE SC'20 paper [Reducing Communicati **Contact:** Alok Tripathy () ## Dependencies -- Python 3.6.10 -- PyTorch 1.3.1 -- PyTorch Geometric (PyG) 1.3.2 -- CUDA 10.1 -- GCC 6.4.0 +- Python 3.8.11 +- PyTorch 1.9.0 +- PyTorch Geometric (PyG) 1.7.2 +- CUDA 10.2 +- GCC 9.3.0 + +On NERSC Perlmutter, all of these dependencies can be accessed with the following +```bash +module load tudatoolkit/21.3_10.2 # CUDA 10.2 +module load pytorch/1.9.0 # PyTorch and PyG +``` -On OLCF Summit, all of these dependencies can be accessed with the following +On OLCF Summit, older versions of these dependencies can be accessed with the following ```bash module load cuda # CUDA 10.1 module load gcc # GCC 6.4.0 @@ -45,45 +49,39 @@ python setup.py install ## Documentation Each algorithm in CAGNET is implemented in a separate file. -- `gcn_distr.py` : 1D algorithm -- `gcn_distr_15d.py` : 1.5D algorithm -- `gcn_distr_2d.py` : 2D algorithm -- `gcn_distr_3d.py` : 3D algorithm +- `examples/gcn_1d.py` : 1D algorithm +- `examples/gcn_15d.py` : 1.5D algorithm +- `examples/gcn_2d.py` : 2D algorithm +- `examples/gcn_3d.py` : 3D algorithm Each file also as the following flags: -- `--accperrank ` : Number of GPUs on each node -- `--epochs ` : Number of epochs to run training -- `--graphname ` : Graph dataset to run training on +- `--n-epochs ` : Number of epochs to run training +- `--dataset ` : Graph dataset to run training on - `--timing ` : Enable timing barriers to time phases in training -- `--midlayer ` : Number of activations in the hidden layer -- `--runcount ` : Number of times to run training -- `--normalization ` : Normalize adjacency matrix in preprocessing -- `--activations ` : Enable activation functions between layers +- `--n-hidden ` : Number of activations in the hidden layer +- `--normalize` : Normalize adjacency matrix in preprocessing - `--accuracy ` : Compute and print accuracy metrics (Reddit only) - `--replication ` : Replication factor (1.5D algorithm only) -- `--download ` : Download the Reddit dataset Some of these flags do not currently exist for the 3D algorithm. -Amazon/Protein datasets must exist as COO files in `../data//processed/`, compressed with pickle. -For Reddit, PyG handles downloading and accessing the dataset (see below). +Amazon/Protein datasets must exist as COO files in `../../data//processed/`, compressed with pickle. +For Cora and Reddit, PyG handles downloading and accessing the dataset. -## Running on OLCF Summit (example) +## Running on NERSC Perlmutter (example) To run the CAGNET 1.5D algorithm on Reddit with -- 16 processes +- 4 processes - 100 epochs - 16 hidden layer activations - 2-factor replication +- normalization +- no weight decay -run the following command to download the Reddit dataset: - -`python gcn_distr_15d.py --graphname=Reddit --download=True` - -This will download Reddit into `../data`. After downloading the Reddit dataset, run the following command to run training +run the following command with Slurm: -`ddlrun -x WORLD_SIZE=16 -x MASTER_ADDR=$(echo $LSB_MCPU_HOSTS | cut -d " " -f 3) -x MASTER_PORT=1234 -accelerators 6 python gcn_distr_15d.py --accperrank=6 --epochs=100 --graphname=Reddit --timing=False --midlayer=16 --runcount=1 --replication=2` +`cd examples && srun -n 4 --ntasks-per-node=4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset=Reddit --n-epochs=100 --n-hidden 16 --weight-decay 0 --replication 2 --normalize` ## Citation diff --git a/cagnet/__init__.py b/cagnet/__init__.py new file mode 100644 index 0000000..efdcbe4 --- /dev/null +++ b/cagnet/__init__.py @@ -0,0 +1 @@ +__all__ = [ "nn", "partitionings" ] diff --git a/cagnet/nn/conv/__init__.py b/cagnet/nn/conv/__init__.py new file mode 100644 index 0000000..73bc91b --- /dev/null +++ b/cagnet/nn/conv/__init__.py @@ -0,0 +1,7 @@ +from .gcn_conv import GCNConv + +__all__ = [ + 'GCNConv', +] + +classes = __all__ diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py new file mode 100644 index 0000000..6efd84b --- /dev/null +++ b/cagnet/nn/conv/gcn_conv.py @@ -0,0 +1,1045 @@ +import math +import torch +import torch.distributed as dist +import torch.nn as nn +import time + +from cagnet.partitionings import Partitioning +from cagnet.samplers.utils import * +from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu, normalize_batch_gpu + +def broad_func_oned(self, graph, ampbyp, inputs): + n_per_proc = math.ceil(float(graph.size(0) / self.size)) + + z_loc = torch.cuda.FloatTensor(ampbyp[0].size(0), inputs.size(1), device=self.device).fill_(0) + + inputs_recv = torch.cuda.FloatTensor(n_per_proc, inputs.size(1), device=self.device).fill_(0) + + for i in range(self.size): + if i == self.rank: + inputs_recv = inputs.clone() + elif i == self.size - 1: + inputs_recv = torch.cuda.FloatTensor(ampbyp[i].size(1), \ + inputs.size(1), \ + device=self.device).fill_(0) + + dist.broadcast(inputs_recv, src=i, group=self.group) + + spmm_gpu(ampbyp[i].indices()[0].int(), ampbyp[i].indices()[1].int(), + ampbyp[i].values(), ampbyp[i].size(0), + ampbyp[i].size(1), inputs_recv, z_loc) + + return z_loc + +def broad_func_one5d(self, graph, ampbyp, inputs): + n_per_proc = math.ceil(float(graph.size(0)) / (self.size / self.replication)) + + z_loc = torch.cuda.FloatTensor(ampbyp[0].size(0), inputs.size(1), device=self.device).fill_(0) + + inputs_recv = torch.cuda.FloatTensor(n_per_proc, inputs.size(1), device=self.device).fill_(0) + + rank_c = self.rank // self.replication + rank_col = self.rank % self.replication + + stages = self.size // (self.replication ** 2) + if rank_col == self.replication - 1: + stages = (self.size // self.replication) - (self.replication - 1) * stages + + for i in range(stages): + q = (rank_col * (self.size // (self.replication ** 2)) + i) * self.replication + rank_col + + q_c = q // self.replication + + am_partid = rank_col * (self.size // self.replication ** 2) + i + + if q == self.rank: + inputs_recv = inputs.clone() + elif q_c == self.size // self.replication - 1: + inputs_recv = torch.cuda.FloatTensor(ampbyp[am_partid].size(1), \ + inputs.size(1), \ + device=self.device).fill_(0) + + inputs_recv = inputs_recv.contiguous() + dist.broadcast(inputs_recv, src=q, group=self.col_groups[rank_col]) + + spmm_gpu(ampbyp[am_partid].indices()[0].int(), ampbyp[am_partid].indices()[1].int(), + ampbyp[am_partid].values(), ampbyp[am_partid].size(0), + ampbyp[am_partid].size(1), inputs_recv, z_loc) + + z_loc = z_loc.contiguous() + + dist.all_reduce(z_loc, op=dist.reduce_op.SUM, group=self.row_groups[rank_c]) + + return z_loc + +def summa_sparse(self, adj_matrix, inputs, height, middim, width): + + height_per_proc = height // self.proc_row + width_per_proc = width // self.proc_col + + middim_per_proc = middim // self.proc_col + + if self.rank_row == self.proc_row - 1: + height_per_proc = height - height_per_proc * (self.proc_row - 1) + + if self.rank_col == self.proc_col - 1: + width_per_proc = width - width_per_proc * (self.proc_col - 1) + + z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=self.device).fill_(0) + + for k in range(self.proc_col): + + row_src_rank = k + self.proc_col * self.rank_row + col_src_rank = k * self.proc_col + self.rank_col + + if k == self.proc_col - 1: + middim_per_proc = middim - middim_per_proc * (self.proc_col - 1) + + if row_src_rank == self.rank: + acol_indices_len = torch.cuda.LongTensor( + [adj_matrix.indices().contiguous()[0].size(0)], + device=self.device) + acol_values_len = torch.cuda.LongTensor([adj_matrix.values().contiguous().size(0)], + device=self.device) + else: + acol_indices_len = torch.cuda.LongTensor([0], device=self.device) + acol_values_len = torch.cuda.LongTensor([0], device=self.device) + + dist.broadcast(acol_indices_len, row_src_rank, self.row_groups[self.rank_row]) + + acol_indices_len = acol_indices_len.item() # nnz + acol_values_len = acol_indices_len + + if row_src_rank == self.rank: + acol_indices = adj_matrix.indices().contiguous().long() + acol_values = adj_matrix.values().contiguous().float() + else: + acol_indices = torch.cuda.LongTensor(2, acol_indices_len, device=self.device).fill_(0) + acol_values = torch.cuda.FloatTensor(acol_values_len, device=self.device).fill_(0) + + acol = torch.cat((acol_indices.float(), acol_values.unsqueeze(0)), dim=0).contiguous() + + dist.broadcast(acol, row_src_rank, self.row_groups[self.rank_row]) + + acol_indices = acol[:2].long() + acol_values = acol[2].squeeze(0) + + if row_src_rank == self.rank: + acol = adj_matrix + else: + acol = sparse_coo_tensor_gpu(acol_indices, acol_values, + torch.Size([height_per_proc, middim_per_proc])) + + if col_src_rank == self.rank: + brow = inputs + else: + brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=self.device) + + brow = brow.contiguous() + + dist.broadcast(brow, col_src_rank, self.col_groups[self.rank_col]) + + spmm_gpu(acol_indices[0].int(), acol_indices[1].int(), acol_values, + height_per_proc, middim_per_proc, brow, z_loc) + + return z_loc + +def summa_loc(self, mata, matb, height, middim, width): + + height_per_proc = height // self.proc_row + width_per_proc = width // self.proc_col + middim_per_proc = middim // self.proc_row + + if self.rank_row == self.proc_row - 1: + height_per_proc = height - height_per_proc * (self.proc_row - 1) + + width_per_proc = matb[self.rank].size(1) + + acol_tens = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=self.device) + brow_tens = torch.FloatTensor(middim_per_proc, width_per_proc) + + acol = acol_tens + brow = brow_tens + + z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=self.device).fill_(0) + + for k in range(self.proc_col): + + row_src_rank = k + self.proc_col * self.rank_row + col_src_rank = k * self.proc_col + self.rank_col + + if k == self.proc_col - 1: + middim_per_proc -= self.proc_col * middim_per_proc - middim + + if row_src_rank == self.rank: + acol = mata + else: + acol = acol_tens + acol = torch.cuda.FloatTensor(height_per_proc, matb[col_src_rank].size(0), + device=self.device) + + acol = acol.contiguous() + dist.broadcast(acol, row_src_rank, self.row_groups[self.rank_row]) + + brow = matb[col_src_rank] + + z_loc += torch.mm(acol, brow) + + return z_loc + +def summa(self, adj_matrix, inputs, height, middim, width): + + height_per_proc = height // self.proc_row + width_per_proc = width // self.proc_col + middim_per_proc = middim // self.proc_row + + if self.rank_row == self.proc_row - 1: + height_per_proc = height - height_per_proc * (self.proc_row - 1) + + if self.rank_col == self.proc_col - 1: + width_per_proc = width - width_per_proc * (self.proc_col - 1) + + acol_tens = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=self.device) + brow_tens = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=self.device) + + acol = acol_tens + brow = brow_tens + + z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=self.device).fill_(0) + + for k in range(self.proc_col): + + row_src_rank = k + self.proc_col * self.rank_row + col_src_rank = k * self.proc_col + self.rank_col + + if k == self.proc_col - 1: + middim_per_proc = middim - middim_per_proc * (self.proc_col - 1) + acol_tens = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=self.device) + brow_tens = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=self.device) + + if row_src_rank == self.rank: + acol = adj_matrix + else: + acol = acol_tens + + acol = acol.contiguous() + dist.broadcast(acol, row_src_rank, self.row_groups[self.rank_row]) + + if col_src_rank == self.rank: + brow = inputs + else: + brow = brow_tens + + brow = brow.contiguous() + dist.broadcast(brow, col_src_rank, self.col_groups[self.rank_col]) + + z_loc += torch.mm(acol.float(), brow) + + return z_loc + +def transpose(self, mat, height, width, height_c=None, width_c=None): + if self.partitioning == Partitioning.TWOD: + return transpose_TWOD(self, mat, height, width) + elif self.partitioning == Partitioning.THREED: + return transpose_THREED(self, mat, height, width, height_c, width_c) + +def transpose_TWOD(self, mat, height, width): + rank = self.rank_row * self.proc_col + self.rank_col + rank_t = self.rank_col * self.proc_row + self.rank_row + + if rank == rank_t: + return mat.t() + + height_recv = width // self.proc_row + width_recv = height // self.proc_col + + if self.rank_row == self.proc_row - 1: + height_recv = width - height_recv * (self.proc_row - 1) + + if self.rank_col == self.proc_col - 1: + width_recv = height - width_recv * (self.proc_col - 1) + + mat_recv = torch.cuda.FloatTensor(height_recv, width_recv, device=self.device) + + mat_recvs = [mat.t().contiguous(), mat_recv] + + if self.rank < rank_t: + dist.broadcast(mat_recvs[0], src=self.rank, group=self.transpose_group) + dist.broadcast(mat_recvs[1], src=rank_t, group=self.transpose_group) + else: + dist.broadcast(mat_recvs[1], src=rank_t, group=self.transpose_group) + dist.broadcast(mat_recvs[0], src=self.rank, group=self.transpose_group) + + return mat_recvs[1] + +def transpose_THREED(self, mat, height, width, height_c, width_c): + rank_t = self.rank_col * self.proc_col * self.proc_c + self.rank_row * self.proc_c + self.rank_c + no_occur_val = 42.1234 + + width_per_proc_c = width_c // self.proc_c + mat_c_recv_width = width_c - width_per_proc_c * (self.proc_c - 1) + + if mat.size(1) != mat_c_recv_width: + pad_col = mat_c_recv_width - mat.size(1) + mat = torch.cat( + (mat, torch.cuda.FloatTensor(mat.size(0), pad_col, device=self.device).fill_(no_occur_val)), + dim=1) + + mat_c_recv = [] + for i in range(self.proc_c): + mat_c_recv.append(torch.cuda.FloatTensor(mat.size(), device=self.device)) + + dist.all_gather(mat_c_recv, mat, group=self.c_groups[int(self.rank // self.proc_c)]) + for i in range(self.proc_c): + mat_c_recv[i] = mat_c_recv[i].t_() + mat_c_recv[i] = mat_c_recv[i][(mat_c_recv[i][:, 0] != no_occur_val).nonzero().squeeze(1)] + mat_c_recv[i] = mat_c_recv[i].t_() + + mat = torch.cat(mat_c_recv, dim=1) + mat.t_() + + if self.rank_col == self.proc_row - 1: + width_recv = height - (height // self.proc_row) * (self.proc_row - 1) + else: + width_recv = height // self.proc_row + + if self.rank_row == self.proc_col - 1: + height_recv = width - (width // self.proc_col) * (self.proc_col - 1) + else: + height_recv = width // self.proc_col + + if self.rank_row == self.rank_col: + mat_recv = mat + else: + mat_recv = torch.cuda.FloatTensor(height_recv, width_recv, device=self.device) + mat_recvs = [mat.contiguous(), mat_recv] + + if self.rank < rank_t: + dist.broadcast(mat_recvs[0], src=self.rank, group=self.transpose_group) + dist.broadcast(mat_recvs[1], src=rank_t, group=self.transpose_group) + else: + dist.broadcast(mat_recvs[1], src=rank_t, group=self.transpose_group) + dist.broadcast(mat_recvs[0], src=self.rank, group=self.transpose_group) + + chunk_sizes_col = [] + chunk_len = width_recv // self.proc_c + for i in range(self.proc_c): + if i == self.proc_c - 1: + chunk_sizes_col.append(width_recv - chunk_len * (self.proc_c - 1)) + else: + chunk_sizes_col.append(chunk_len) + mat_split = torch.split(mat_recv, chunk_sizes_col, dim=1) + + return mat_split[self.rank_c] + +def outer_product(mata, matb, group): + matc = torch.mm(mata, matb) + dist.all_reduce(matc, op=dist.reduce_op.SUM, group=group) + + return matc + +def split3dspmm_sparse(self, adj_matrix, inputs, height, middim, width): + + height_per_proc = height // self.proc_row + width_per_proc = width // self.proc_col + middim_per_proc = middim // (self.proc_col * self.proc_c) + + if self.rank_row == self.proc_row - 1: + height_per_proc = height - height_per_proc * (self.proc_row - 1) + + if self.rank_col == self.proc_col - 1: + width_per_proc = width - width_per_proc * (self.proc_col - 1) + + width_per_proc_c = inputs.size(1) // self.proc_c + if self.rank_c == self.proc_c - 1: + width_per_proc_c = inputs.size(1) - width_per_proc_c * (self.proc_c - 1) + + z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=self.device).fill_(0) + + chunk_sizes_col = [] + chunk_len = inputs.size(1) // self.proc_c + for i in range(self.proc_c): + if i == self.proc_c - 1: + chunk_sizes_col.append(inputs.size(1) - chunk_len * (self.proc_c - 1)) + else: + chunk_sizes_col.append(chunk_len) + + for k in range(self.proc_col): + + row_src_rank = self.rank_row * (self.proc_col * self.proc_c) + self.rank_c + k * self.proc_c + col_src_rank = self.rank_col * self.proc_row + self.rank_c + k * self.proc_c * self.proc_row + + middim_per_col = middim // self.proc_col + if k == self.proc_col - 1: + middim_per_col = middim - middim_per_col * (self.proc_col - 1) + + middim_per_proc = middim_per_col // self.proc_c + if self.rank_c == self.proc_c - 1: + middim_per_proc = middim_per_col - middim_per_proc * (self.proc_c - 1) + + if row_src_rank == self.rank: + acol_indices_len = torch.cuda.LongTensor( + [adj_matrix.indices().contiguous()[0].size(0)], + device=self.device) + acol_values_len = torch.cuda.LongTensor([adj_matrix.values().contiguous().size(0)], + device=self.device) + else: + acol_indices_len = torch.cuda.LongTensor([0], device=self.device) + acol_values_len = torch.cuda.LongTensor([0], device=self.device) + + dist.broadcast(acol_indices_len, row_src_rank, self.row_groups[self.rank_row][self.rank_c]) + + acol_indices_len = acol_indices_len.item() # nnz + acol_values_len = acol_indices_len + + if row_src_rank == self.rank: + acol_indices = adj_matrix.indices().contiguous().long() + acol_values = adj_matrix.values().contiguous().float() + else: + acol_indices = torch.cuda.LongTensor(2, acol_indices_len, device=self.device).fill_(0) + acol_values = torch.cuda.FloatTensor(acol_values_len, device=self.device).fill_(0) + + acol = torch.cat((acol_indices.float(), acol_values.unsqueeze(0)), dim=0) + + dist.broadcast(acol.contiguous(), row_src_rank, self.row_groups[self.rank_row][self.rank_c]) + + acol_indices = acol[:2].long() + acol_values = acol[2].squeeze(0) + + if row_src_rank == self.rank: + acol = adj_matrix + else: + acol = sparse_coo_tensor_gpu(acol_indices, acol_values, + torch.Size([height_per_proc, middim_per_proc])) + + if col_src_rank == self.rank: + brow = inputs + else: + brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=self.device) + + brow = brow.contiguous() + dist.broadcast(brow, col_src_rank, self.col_groups[self.rank_col][self.rank_c]) + + spmm_gpu(acol_indices[0].int(), acol_indices[1].int(), acol_values, + height_per_proc, middim_per_proc, brow, z_loc) + + z_loc = z_loc.contiguous() + + dist.all_reduce(z_loc, group=self.c_groups[int(self.rank // self.proc_c)]) + z_loc = torch.split(z_loc, chunk_sizes_col, dim=1) + z_loc = z_loc[self.rank_c].contiguous() + + return z_loc, chunk_sizes_col + +def split3dspmm_loc(self, mata, matb, height, middim, width): + + height_per_proc = height // self.proc_row + width_per_proc = width // self.proc_col + middim_per_proc = middim // (self.proc_col * self.proc_c) + + if self.rank_row == self.proc_row - 1: + height_per_proc = height - height_per_proc * (self.proc_row - 1) + + width_per_proc = matb[self.rank].size(1) + + height_per_proc_c = mata.size(0) // self.proc_c + if self.rank_c == self.proc_c - 1: + height_per_proc_c = mata.size(0) - height_per_proc_c * (self.proc_c - 1) + + z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=self.device).fill_(0) + + chunk_sizes_row = [] + chunk_len = mata.size(0) // self.proc_c + for i in range(self.proc_c): + if i == self.proc_c - 1: + chunk_sizes_row.append(mata.size(0) - chunk_len * (self.proc_c - 1)) + else: + chunk_sizes_row.append(chunk_len) + + for k in range(self.proc_col): + + row_src_rank = self.rank_row * (self.proc_col * self.proc_c) + self.rank_c + k * self.proc_c + col_src_rank = self.rank_col * self.proc_row + self.rank_c + k * self.proc_c * self.proc_row + + middim_per_col = middim // self.proc_col + if k == self.proc_col - 1: + middim_per_col = middim - middim_per_col * (self.proc_col - 1) + + middim_per_proc = middim_per_col // self.proc_c + if self.rank_c == self.proc_c - 1: + middim_per_proc = middim_per_col - middim_per_proc * (self.proc_c - 1) + + if row_src_rank == self.rank: + acol = mata + else: + acol = torch.cuda.FloatTensor(height_per_proc, matb[col_src_rank].size(0), + device=self.device) + + dist.broadcast(acol.contiguous(), row_src_rank, self.row_groups[self.rank_row][self.rank_c]) + + brow = matb[col_src_rank] + + z_tmp = torch.mm(acol, brow) + + dist.all_reduce(z_tmp, group=self.c_groups[int(self.rank // self.proc_c)]) + + z_loc += z_tmp + + z_tmp = torch.split(z_loc, chunk_sizes_row, dim=0) + z_loc = z_tmp[self.rank_c].clone() + return z_loc + +def split3dspmm_dense(self, adj_matrix, inputs, height, middim, width): + + height_per_proc = height // self.proc_row + width_per_proc = width // self.proc_col + middim_per_proc = middim // (self.proc_col * self.proc_c) + + if self.rank_row == self.proc_row - 1: + height_per_proc = height - height_per_proc * (self.proc_row - 1) + + if self.rank_col == self.proc_col - 1: + width_per_proc = width - width_per_proc * (self.proc_col - 1) + + height_per_proc_c = adj_matrix.size(0) // self.proc_c + if self.rank_c == self.proc_c - 1: + height_per_proc_c = adj_matrix.size(0) - height_per_proc_c * (self.proc_c - 1) + + z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=self.device).fill_(0) + + chunk_sizes_row = [] + chunk_len = adj_matrix.size(0) // self.proc_c + for i in range(self.proc_c): + if i == self.proc_c - 1: + chunk_sizes_row.append(adj_matrix.size(0) - chunk_len * (self.proc_c - 1)) + else: + chunk_sizes_row.append(chunk_len) + + for k in range(self.proc_col): + + row_src_rank = self.rank_row * (self.proc_col * self.proc_c) + self.rank_c + k * self.proc_c + col_src_rank = self.rank_col * self.proc_row + self.rank_c + k * self.proc_c * self.proc_row + + middim_per_col = middim // self.proc_col + if k == self.proc_col - 1: + middim_per_col = middim - middim_per_col * (self.proc_col - 1) + + middim_per_proc = middim_per_col // self.proc_c + if self.rank_c == self.proc_c - 1: + middim_per_proc = middim_per_col - middim_per_proc * (self.proc_c - 1) + + if row_src_rank == self.rank: + acol = adj_matrix + else: + acol = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=self.device).fill_(0) + + acol = acol.contiguous() + dist.broadcast(acol, row_src_rank, self.row_groups[self.rank_row][self.rank_c]) + + if col_src_rank == self.rank: + brow = inputs + else: + brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=self.device).fill_(0) + + brow = brow.contiguous() + dist.broadcast(brow, col_src_rank, self.col_groups[self.rank_col][self.rank_c]) + + z_loc += torch.mm(acol, brow) + + dist.all_reduce(z_loc, group=self.c_groups[int(self.rank // self.proc_c)]) + z_loc = torch.split(z_loc, chunk_sizes_row, dim=0) + z_loc = z_loc[self.rank_c].contiguous() + + return z_loc + +class GCNConv(nn.Module): + def __init__(self, in_feats, out_feats, partitioning, device, aggr="sum"): + super(GCNConv, self).__init__() + + weight_nonleaf = torch.rand(in_feats, out_feats, requires_grad=True) + weight_nonleaf = weight_nonleaf.to(device) + weight_nonleaf.retain_grad() + self.weight = nn.Parameter(weight_nonleaf) + self.partitioning = partitioning + self.aggr = aggr + + def forward(self, gcn, graph, inputs, epoch=0, ampbyp=None): + if self.partitioning == Partitioning.ONED: + return GCNFuncONED.apply(gcn, graph, ampbyp, inputs, self.weight) + elif self.partitioning == Partitioning.ONE5D: + return GCNFuncONE5D.apply(gcn, graph, ampbyp, inputs, self.weight) + elif self.partitioning == Partitioning.TWOD: + return GCNFuncTWOD.apply(gcn, graph, inputs, self.weight) + elif self.partitioning == Partitioning.THREED: + return GCNFuncTHREED.apply(gcn, graph, inputs, self.weight) + elif self.partitioning == Partitioning.NONE: + return GCNFuncNONE.apply(gcn, graph, inputs, self.weight, epoch) + + +class GCNFuncNONE(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, inputs, weight, epoch=0): + # inputs: H + # graph: A + # weight: W + + if epoch == -1: # inference + z = torch.FloatTensor(graph.size(0), inputs.size(1)).fill_(0) + z += torch.sparse.mm(graph, inputs) + z = z.mm(weight.cpu()) + if self.aggr == "mean": + degs = torch.IntTensor(graph.size(0)).fill_(0) + ones = torch.IntTensor(graph._indices()[0,:].size(0)).fill_(1) + degs.scatter_add_(0, graph._indices()[0,:], ones) + degs.clamp_(min=1) + z = torch.div(z, degs.unsqueeze(1)) + return z + + z = torch.cuda.FloatTensor(graph.size(0), inputs.size(1), \ + device=self.device).fill_(0) + + spmm_gpu(graph.indices()[0].int(), graph.indices()[1].int(), + graph.values(), graph.size(0), + graph.size(1), inputs, z) + + z = z.mm(weight) + + if self.aggr == "mean": + degs = torch.cuda.IntTensor(graph.size(0)).fill_(0) + ones = torch.cuda.IntTensor(graph._indices()[0,:].size(0)).fill_(1) + degs.scatter_add_(0, graph._indices()[0,:], ones) + degs.clamp_(min=1) + normalize_batch_gpu(z, degs, z.size(0), z.size(1)) + + ctx.save_for_backward(inputs, weight) + ctx.graph = graph + # ctx.self = self + ctx.device = self.device + ctx.aggr = self.aggr + ctx.epoch = epoch + ctx.timings = self.timings + + return z + + @staticmethod + def backward(ctx, grad_output): + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + epoch = ctx.epoch + if epoch >= 1: + start_time(start_timer) + graph = ctx.graph + inputs, weight = ctx.saved_tensors + # self = ctx.self + device = ctx.device + aggr = ctx.aggr + timings = ctx.timings + + ag = torch.cuda.FloatTensor(graph.size(1), grad_output.size(1), \ + device=device).fill_(0) + if epoch >= 1: + timings["precomp"].append(stop_time(start_timer, stop_timer)) + + # Use A^T instead of A implicitly + if epoch >= 1: + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-bwd-spmm") + spmm_gpu(graph.indices()[1].int(), graph.indices()[0].int(), + graph.values(), graph.size(1), + graph.size(0), grad_output, ag) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + timings["spmm"].append(stop_time(start_timer, stop_timer)) + + if epoch >= 1: + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-bwd-aggr") + if aggr == "mean": + degs = torch.cuda.IntTensor(graph.size(1)).fill_(0) + ones = torch.cuda.IntTensor(graph._indices()[0,:].size(0)).fill_(1) + degs.scatter_add_(0, graph._indices()[1,:], ones) + degs.clamp_(min=1) + normalize_batch_gpu(ag, degs, ag.size(0), ag.size(1)) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + timings["aggr"].append(stop_time(start_timer, stop_timer)) + + if epoch >= 1: + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-bwd-gemmi") + grad_input = ag.mm(weight.t()) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + timings["gemm_i"].append(stop_time(start_timer, stop_timer)) + + if epoch >= 1: + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-bwd-gemmw") + grad_weight = inputs.t().mm(ag) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + timings["gemm_w"].append(stop_time(start_timer, stop_timer)) + + return None, None, grad_input, grad_weight, None + +class GCNFuncONED(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, ampbyp, inputs, weight): + # inputs: H + # graph: A + # weight: W + + z = broad_func_oned(self, graph, ampbyp, inputs) + z = z.mm(weight) + + ctx.save_for_backward(inputs, weight) + ctx.ampbyp = ampbyp + ctx.graph = graph + ctx.self = self + + return z + + @staticmethod + def backward(ctx, grad_output): + graph = ctx.graph + ampbyp = ctx.ampbyp + inputs, weight = ctx.saved_tensors + self = ctx.self + + # Assumes graph is undirected and A = A^T + ag = broad_func_oned(self, graph, ampbyp, grad_output) + + grad_input = ag.mm(weight.t()) + grad_weight = outer_product(inputs.t(), ag, self.group) + + return None, None, None, grad_input, grad_weight + +class GCNFuncONE5D(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, ampbyp, inputs, weight): + # inputs: H + # graph: A + # weight: W + + z = broad_func_one5d(self, graph, ampbyp, inputs) + z = z.mm(weight) + + ctx.save_for_backward(inputs, weight) + ctx.ampbyp = ampbyp + ctx.graph = graph + ctx.self = self + + return z + + @staticmethod + def backward(ctx, grad_output): + graph = ctx.graph + ampbyp = ctx.ampbyp + inputs, weight = ctx.saved_tensors + self = ctx.self + + # Assumes graph is undirected and A = A^T + ag = broad_func_one5d(self, graph, ampbyp, grad_output) + + grad_input = ag.mm(weight.t()) + grad_weight = outer_product(inputs.t(), ag, self.group) + + return None, None, None, grad_input, grad_weight + +class GCNFuncTWOD(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, inputs, weight): + # inputs: H + # graph: A + # weight: W + + graph_t = graph # Only true for undirected graphs + + z = summa_sparse(self, graph_t, inputs, self.node_count, self.node_count, weight.size(0)) + + chunk_sizes_row = [] + chunk_sizes_col = [] + weight_per_row = weight.size(0) // self.proc_row + weight_per_col = weight.size(1) // self.proc_col + for i in range(self.proc_row): + if i == self.proc_row - 1: + chunk_sizes_row.append(weight.size(0) - weight_per_row * (self.proc_row - 1)) + else: + chunk_sizes_row.append(weight_per_row) + + for i in range(self.proc_col): + if i == self.proc_col - 1: + chunk_sizes_col.append(weight.size(1) - weight_per_col * (self.proc_col - 1)) + else: + chunk_sizes_col.append(weight_per_col) + + weight_rows = torch.split(weight, chunk_sizes_row, dim=0) + weight_parts = [] + for i in weight_rows: + weight_cols = torch.split(i, chunk_sizes_col, dim=1) + weight_parts.extend(weight_cols) + + z = summa_loc(self, z, weight_parts, self.node_count, weight.size(0), weight.size(1)) + + ctx.save_for_backward(inputs, weight) + ctx.graph = graph + ctx.self = self + + return z + + @staticmethod + def backward(ctx, grad_output): + graph = ctx.graph + inputs, weight = ctx.saved_tensors + self = ctx.self + + # Assumes graph is undirected and A = A^T + ag = summa_sparse(self, graph, grad_output, self.node_count, self.node_count, weight.t().size(0)) + + chunk_sizes_row = [] + chunk_sizes_col = [] + weight_per_row = weight.t().size(0) // self.proc_row + weight_per_col = weight.t().size(1) // self.proc_col + + for i in range(self.proc_row): + if i == self.proc_row - 1: + chunk_sizes_row.append(weight.t().size(0) - weight_per_row * (self.proc_row - 1)) + else: + chunk_sizes_row.append(weight_per_row) + + for i in range(self.proc_col): + if i == self.proc_col - 1: + chunk_sizes_col.append(weight.t().size(1) - weight_per_col * (self.proc_col - 1)) + else: + chunk_sizes_col.append(weight_per_col) + + weight_rows = torch.split(weight.t(), chunk_sizes_row, dim=0) + + weight_parts = [] + for i in weight_rows: + weight_cols = torch.split(i, chunk_sizes_col, dim=1) + weight_parts.extend(weight_cols) + + grad_input = summa_loc(self, ag, weight_parts, self.node_count, weight.t().size(0), weight.t().size(1)) + + # Second backprop equation (reuses the A * G^l computation) + inputs_t = transpose(self, inputs, self.node_count, weight.size(0)) + + grad_weight = summa(self, inputs_t, ag, weight.size(0), self.node_count, weight.size(1)) + + # Collect grad_weight's across processes + grad_weight_recv = [] + max_row_chunk = max(chunk_sizes_col) # transpose + max_col_chunk = max(chunk_sizes_row) + for i in range(self.size): + grad_weight_recv.append(torch.cuda.FloatTensor( + max_row_chunk, + max_col_chunk, + device=self.device)) + + pad_row = max_row_chunk - grad_weight.size(0) + pad_col = max_col_chunk - grad_weight.size(1) + + # TODO: make this part less hacky + no_occur_val = 42.1234 + grad_weight = torch.cat((grad_weight, + torch.cuda.FloatTensor(pad_row, grad_weight.size(1), device=self.device).fill_(no_occur_val)), + dim=0) + grad_weight = torch.cat((grad_weight, + torch.cuda.FloatTensor(grad_weight.size(0), pad_col, device=self.device).fill_(no_occur_val)), + dim=1) + + dist.all_gather(grad_weight_recv, grad_weight) + + for i in range(len(grad_weight_recv)): + grad_weight_recv[i] = grad_weight_recv[i][(grad_weight_recv[i][:, 0] != no_occur_val) + .nonzero().squeeze(1)] + + grad_weight_recv_t = grad_weight_recv[i].t() + grad_weight_recv_t = grad_weight_recv_t[(grad_weight_recv_t[:, 0] != no_occur_val) + .nonzero().squeeze(1)] + + grad_weight_recv[i] = grad_weight_recv_t.t() + + grad_weight_fin = torch.cuda.FloatTensor(device=self.device) + for i in range(self.proc_row): + grad_weight_row = torch.cuda.FloatTensor(device=self.device) + for j in range(self.proc_col): + rank_wt = i * self.proc_row + j + grad_weight_row = torch.cat((grad_weight_row, grad_weight_recv[rank_wt]), dim=1) + grad_weight_fin = torch.cat((grad_weight_fin, grad_weight_row), dim=0) + + return None, None, grad_input, grad_weight_fin + +class GCNFuncTHREED(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, inputs, weight): + # inputs: H + # graph: A + # weight: W + + graph_t = graph # Only true for undirected graphs + + z, chunk_sizes_loc = split3dspmm_sparse(self, graph_t, inputs, self.node_count, + self.node_count, weight.size(0)) + + chunk_sizes_loc_tens = torch.cuda.LongTensor(chunk_sizes_loc) + chunk_sizes = [] + for i in range(self.proc_col): + chunk_sizes.append(torch.cuda.LongTensor(chunk_sizes_loc_tens.size())) + dist.all_gather(chunk_sizes, chunk_sizes_loc_tens, group=self.row_groups[self.rank_row][self.rank_c]) + chunk_sizes = torch.cat(chunk_sizes).tolist() + + chunk_sizes_row = [] + chunk_sizes_col = [] + weight_per_row = weight.size(0) // (self.proc_row * self.proc_c) + weight_per_col = weight.size(1) // self.proc_col + chunk_sizes_row = chunk_sizes + + for i in range(self.proc_col): + if i == self.proc_col - 1: + chunk_sizes_col.append(weight.size(1) - weight_per_col * (self.proc_col - 1)) + else: + chunk_sizes_col.append(weight_per_col) + + weight_rows = torch.split(weight, chunk_sizes_row, dim=0) + weight_parts_tmp = [] + for i in weight_rows: + weight_cols = torch.split(i, chunk_sizes_col, dim=1) + weight_parts_tmp.extend(weight_cols) + + weight_parts = [None] * self.size + for i in range(self.proc_row * self.proc_c): + for j in range(self.proc_col): + rank_old = i * self.proc_col + j + + rank_new_row = i // self.proc_c + rank_new_c = i % self.proc_c + rank_new = rank_new_row * self.proc_col * self.proc_c + j * self.proc_c + rank_new_c + + if j == self.rank_col: + weight_parts[rank_new] = weight_parts_tmp[rank_old] + else: + weight_parts[rank_new] = None + + z = split3dspmm_loc(self, z, weight_parts, self.node_count, weight.size(0), weight.size(1)) + + ctx.save_for_backward(inputs, weight) + ctx.graph = graph + ctx.self = self + + return z + + @staticmethod + def backward(ctx, grad_output): + graph = ctx.graph + inputs, weight = ctx.saved_tensors + self = ctx.self + + # First backprop equation + ag, chunk_sizes_loc = split3dspmm_sparse(self, graph, grad_output, self.node_count, + self.node_count, weight.t().size(0)) + + chunk_sizes_loc_tens = torch.cuda.LongTensor(chunk_sizes_loc) + chunk_sizes = [] + for i in range(self.proc_col): + chunk_sizes.append(torch.cuda.LongTensor(chunk_sizes_loc_tens.size())) + dist.all_gather(chunk_sizes, chunk_sizes_loc_tens, group=self.row_groups[self.rank_row][self.rank_c]) + chunk_sizes = torch.cat(chunk_sizes).tolist() + + chunk_sizes_row = [] + chunk_sizes_col = [] + weight_per_row = weight.t().size(0) // (self.proc_row * self.proc_c) + weight_per_col = weight.t().size(1) // self.proc_col + chunk_sizes_row = chunk_sizes + + for i in range(self.proc_col): + if i == self.proc_col - 1: + chunk_sizes_col.append(weight.t().size(1) - weight_per_col * (self.proc_col - 1)) + else: + chunk_sizes_col.append(weight_per_col) + weight_rows = torch.split(weight.t(), chunk_sizes_row, dim=0) + + weight_parts_tmp = [] + for i in weight_rows: + weight_cols = torch.split(i, chunk_sizes_col, dim=1) + weight_parts_tmp.extend(weight_cols) + + weight_parts = [None] * self.size + for i in range(self.proc_row * self.proc_c): + for j in range(self.proc_col): + rank_old = i * self.proc_col + j + + rank_new_row = i // self.proc_c + rank_new_c = i % self.proc_c + rank_new = rank_new_row * self.proc_col * self.proc_c + j * self.proc_c + rank_new_c + + if self.rank_col == j: + weight_parts[rank_new] = weight_parts_tmp[rank_old] + else: + weight_parts[rank_new] = None + + grad_input = split3dspmm_loc(self, ag, weight_parts, self.node_count, weight.t().size(0), weight.size(1)) + + # Second backprop equation (reuses the A * G^l computation) + # col_groups twice because of transpose + height_c = self.node_count // self.proc_row + if self.rank_row == self.proc_row - 1: + height_c = self.node_count - height_c * (self.proc_row - 1) + + width_c = weight.size(1) // self.proc_col + if self.rank_col == self.proc_col - 1: + width_c = weight.size(1) - width_c * (self.proc_col - 1) + + ag_t = transpose(self, ag, self.node_count, weight.size(1), height_c, width_c) + + grad_weight = split3dspmm_dense(self, ag_t, inputs, weight.size(1), self.node_count, weight.size(0)) + + grad_weight_recv = [] + max_row_chunk = max(chunk_sizes_row) + max_col_chunk = max(chunk_sizes_col) + for i in range(self.size): + grad_weight_recv.append(torch.cuda.FloatTensor( + max_row_chunk, + max_col_chunk, + device=self.device)) + + pad_row = max_row_chunk - grad_weight.size(0) + pad_col = max_col_chunk - grad_weight.size(1) + + # TODO: make this part less hacky + no_occur_val = 42.1234 + grad_weight = torch.cat((grad_weight, + torch.cuda.FloatTensor(pad_row, grad_weight.size(1), device=self.device).fill_(no_occur_val)), + dim=0) + grad_weight = torch.cat((grad_weight, + torch.cuda.FloatTensor(grad_weight.size(0), pad_col, device=self.device).fill_(no_occur_val)), + dim=1) + + dist.all_gather(grad_weight_recv, grad_weight) + + for i in range(len(grad_weight_recv)): + grad_weight_recv[i] = grad_weight_recv[i][(grad_weight_recv[i][:, 0] != no_occur_val) + .nonzero().squeeze(1)] + + grad_weight_recv_t = grad_weight_recv[i].t() + grad_weight_recv_t = grad_weight_recv_t[(grad_weight_recv_t[:, 0] != no_occur_val) + .nonzero().squeeze(1)] + + grad_weight_recv[i] = grad_weight_recv_t.t() + + grad_weight_fin = torch.cuda.FloatTensor(device=self.device) + for i in range(self.proc_row): + grad_weight_row = torch.cuda.FloatTensor(device=self.device) + for j in range(self.proc_col): + grad_weight_col = torch.cuda.FloatTensor(device=self.device) + for k in range(self.proc_c): + rank_wt = i * self.proc_row * self.proc_c + j * self.proc_c + k + grad_weight_col = torch.cat((grad_weight_col, grad_weight_recv[rank_wt]), dim=0) + grad_weight_row = torch.cat((grad_weight_row, grad_weight_col), dim=1) + grad_weight_fin = torch.cat((grad_weight_fin, grad_weight_row), dim=0) + + grad_weight_fin = grad_weight_fin.t() + + return None, None, grad_input, grad_weight_fin diff --git a/cagnet/nn/functional.py b/cagnet/nn/functional.py new file mode 100644 index 0000000..a726742 --- /dev/null +++ b/cagnet/nn/functional.py @@ -0,0 +1,273 @@ +from cagnet.partitionings import Partitioning +import math +import torch +import torch.nn.functional as F +import torch.distributed as dist + +def proc_row_size(size, partitioning): + if partitioning == Partitioning.TWOD: + return math.floor(math.sqrt(size)) + elif partitioning == Partitioning.THREED: + cube_root = int(size ** (1./ 3.)) + if cube_root ** 3 == size: + return cube_root + elif (cube_root + 1) ** 3 == size: + return cube_root + 1 + else: + print(f"CUBE ROOT ERROR") + +def proc_col_size(size, partitioning): + if partitioning == Partitioning.TWOD: + return math.floor(math.sqrt(size)) + elif partitioning == Partitioning.THREED: + cube_root = int(size ** (1./ 3.)) + if cube_root ** 3 == size: + return cube_root + elif (cube_root + 1) ** 3 == size: + return cube_root + 1 + else: + print(f"CUBE ROOT ERROR") + +def proc_c_size(size, partitioning): + if partitioning == Partitioning.THREED: + cube_root = int(size ** (1./ 3.)) + if cube_root ** 3 == size: + return cube_root + elif (cube_root + 1) ** 3 == size: + return cube_root + 1 + else: + print(f"CUBE ROOT ERROR") + +def relu(x, partitioning): + return F.relu(x) + +def log_softmax(self, x, partitioning, dim=1): + if partitioning == Partitioning.ONED \ + or partitioning == Partitioning.ONE5D \ + or partitioning == Partitioning.NONE: + + return F.log_softmax(x, dim) + elif partitioning == Partitioning.TWOD: + return LogSoftmaxTWOD.apply(self, x) + +def cross_entropy(logits_rank, labels_rank, node_count, total_classes, partitioning, rank, group, size, \ + train_mask=None): + + if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: + loss = F.cross_entropy(logits_rank, labels_rank, reduction="sum") + + loss_recv = [] + for i in range(size): + loss_recv.append(torch.cuda.FloatTensor(loss.size())) + dist.all_gather(loss_recv, loss, group) + loss_recv[rank] = loss + loss = sum(loss_recv) / node_count + return loss + + elif partitioning == Partitioning.TWOD: + proc_row = proc_row_size(size, partitioning) + proc_col = proc_col_size(size, partitioning) + + rank_row = int(rank / proc_col) + rank_col = rank % proc_col + + class_per_rank = total_classes // proc_col + + row_groups = group[0] + col_groups = group[1] + + min_class = rank_col * class_per_rank + max_class = min((rank_col + 1) * class_per_rank, total_classes) + if rank_col == proc_col - 1: + max_classes = total_classes + + if list(labels_rank.size())[0] > 0: + datay_ids = labels_rank.long() + + filtered_indices = torch.mul(datay_ids >= min_class, datay_ids < max_class).float() + indices = torch.nonzero(filtered_indices * \ + torch.cuda.FloatTensor(datay_ids.size()).fill_(1)).squeeze() + + datay_ids = labels_rank.long().view(-1, 1) + datay_ids = datay_ids.index_select(0, indices) + datay_ids -= min_class + outputs_ids = logits_rank.index_select(0, indices) + + classes = torch.gather(outputs_ids, 1, datay_ids) + loss_calc = torch.sum(classes) + loss_calc_tens = torch.Tensor([loss_calc.item()]) + + rank_row_src = rank_row * proc_col + + dist.reduce(loss_calc, dst=rank_row_src, op=dist.reduce_op.SUM, group=row_groups[rank_row]) + dist.broadcast(loss_calc, src=rank_row_src, group=row_groups[rank_row]) + + vertex_train_count = (train_mask.size(0) - (train_mask == 0).sum(dim=0)) + loss_calc = -loss_calc / vertex_train_count + return loss_calc + + else: + fake_loss = (logits_rank * torch.cuda.FloatTensor(logits_rank.size()).fill_(0)).sum() + return fake_loss + + elif partitioning == Partitioning.THREED: + proc_row = proc_row_size(size, partitioning) + proc_col = proc_col_size(size, partitioning) + proc_c = proc_c_size(size, partitioning) + + rank_row = int((rank // proc_c) // proc_col) # i in process grid + rank_col = int((rank // proc_c) % proc_col) # j in process grid + rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid + + class_per_rank = total_classes // proc_col + + row_groups = group[0] + col_groups = group[1] + c_groups = group[2] + + min_class = rank_col * class_per_rank + max_class = min((rank_col + 1) * class_per_rank, total_classes) + if rank_col == proc_col - 1: + max_class = total_classes + + # Note: bool type removes warnings, unsure of perf penalty + if list(labels_rank.size())[0] > 0: + datay_ids = labels_rank.long() + + filtered_indices = torch.mul(datay_ids >= min_class, datay_ids < max_class).float() + indices = torch.nonzero(filtered_indices * \ + torch.cuda.FloatTensor(datay_ids.size()).fill_(1)).squeeze() + + datay_ids = labels_rank.long().view(-1, 1) + datay_ids = datay_ids.index_select(0, indices) + datay_ids -= min_class + outputs_ids = logits_rank.index_select(0, indices) + + classes = torch.gather(outputs_ids, 1, datay_ids) + loss_calc = torch.sum(classes) + loss_calc_tens = torch.Tensor([loss_calc.item()]) + + dist.all_reduce(loss_calc, op=dist.reduce_op.SUM, group=row_groups[rank_row][rank_c]) + + vertex_train_count = (train_mask.size(0) - (train_mask == 0).sum(dim=0)) + loss_calc = -loss_calc / vertex_train_count + + return loss_calc + else: + fake_loss = (logits_rank * torch.cuda.FloatTensor(logits_rank.size()).fill_(0)).sum() + return fake_loss + + +class LogSoftmaxTWOD(torch.autograd.Function): + @staticmethod + def forward(ctx, self, z): + maxes = torch.max(z, dim=1, keepdim=True)[0] + maxes_recv = [] + for i in range(self.proc_col): + maxes_recv.append(torch.cuda.FloatTensor(maxes.size(), device=self.device)) + + dist.all_gather(maxes_recv, maxes, group=self.row_groups[self.rank_row]) + maxes_recv[self.rank_col] = maxes + maxes = torch.max(torch.cat(maxes_recv, dim=1), dim=1, keepdim=True)[0] + + h = torch.exp(z - maxes) + sm_sum = torch.sum(h, dim=1, keepdim=True) + + sm_sum_recv = [] + for i in range(self.proc_col): + sm_sum_recv.append(torch.cuda.FloatTensor(sm_sum.size(), device=self.device)) + + dist.all_gather(sm_sum_recv, sm_sum, group=self.row_groups[self.rank_row]) + sm_sum_recv[self.rank_col] = sm_sum + sm_sum = torch.sum(torch.cat(sm_sum_recv, dim=1), dim=1, keepdim=True) + sm_sum = torch.log(sm_sum) + h = z - maxes - sm_sum + + ctx.self = self + ctx.z = z + return h + + @staticmethod + def backward(ctx, grad_output): + self = ctx.self + z = ctx.z + + width = self.n_classes + + with torch.set_grad_enabled(True): + chunk_sizes_col = [] + width_per_col = width // self.proc_col + + for i in range(self.proc_col): + if i == self.proc_col - 1: + chunk_sizes_col.append(width - width_per_col * (self.proc_col - 1)) + else: + chunk_sizes_col.append(width_per_col) + + width_per_proc = width - width_per_col * (self.proc_col - 1) + if z.size(1) != width_per_proc: + z = torch.cat((z, torch.cuda.FloatTensor(z.size(0), width_per_proc - z.size(1))), dim=1) + + z_recv = [] + for i in range(self.proc_col): + z_recv.append(torch.cuda.FloatTensor(z.size())) + + dist.all_gather(z_recv, z, group=self.row_groups[self.rank_row]) + z_recv[self.rank_col] = z + + for i in range(self.proc_col - 1): + pad_col = width // self.proc_col + z_recv[i] = z_recv[i][:,:pad_col] + + z = torch.cat(z_recv, dim=1) + + if grad_output is not None: + if grad_output.size(1) != width_per_proc: + grad_output = torch.cat((grad_output, + torch.cuda.FloatTensor(grad_output.size(0), + width_per_proc - grad_output.size(1))), + dim=1) + + grad_output_recv = [] + for i in range(self.proc_col): + grad_output_recv.append(torch.cuda.FloatTensor(grad_output.size())) + + dist.all_gather(grad_output_recv, grad_output, group=self.row_groups[self.rank_row]) + grad_output_recv[self.rank_col] = grad_output + + for i in range(self.proc_col - 1): + pad_col = width // self.proc_col + grad_output_recv[i] = grad_output_recv[i][:,:pad_col] + + grad_output = torch.cat(grad_output_recv, dim=1) + + maxes = torch.max(z, dim=1, keepdim=True)[0] + h = torch.exp(z - maxes) + sm_sum = torch.sum(h, dim=1, keepdim=True) + sm_sum = torch.log(sm_sum) + + h = z - maxes - sm_sum + + func_eval = h + z_gathered = z + grad_output_gathered = grad_output + + width = z_gathered.size(1) + + sigmap = torch.autograd.grad(outputs=func_eval, \ + inputs=z_gathered, \ + grad_outputs=grad_output_gathered)[0] + + chunk_sizes_col = [] + sigmap_per_col = width // self.proc_col + + for i in range(self.proc_col): + if i == self.proc_col - 1: + chunk_sizes_col.append(width - sigmap_per_col * (self.proc_col - 1)) + else: + chunk_sizes_col.append(sigmap_per_col) + + z_grad = sigmap.split(chunk_sizes_col, dim=1)[self.rank_col] + + return None, z_grad + diff --git a/cagnet/partitionings/__init__.py b/cagnet/partitionings/__init__.py new file mode 100644 index 0000000..5ed5bec --- /dev/null +++ b/cagnet/partitionings/__init__.py @@ -0,0 +1,7 @@ +from .partitionings import Partitioning + +__all__ = [ + 'Partitioning', +] + +classes = __all__ diff --git a/cagnet/partitionings/partitionings.py b/cagnet/partitionings/partitionings.py new file mode 100644 index 0000000..fe40743 --- /dev/null +++ b/cagnet/partitionings/partitionings.py @@ -0,0 +1,8 @@ +import enum + +class Partitioning(enum.IntEnum): + ONED = 0 + ONE5D = 1 + TWOD = 2 + THREED = 3 + NONE = 4 diff --git a/cagnet/samplers/__init__.py b/cagnet/samplers/__init__.py new file mode 100644 index 0000000..6f3d75e --- /dev/null +++ b/cagnet/samplers/__init__.py @@ -0,0 +1,11 @@ +from .ladies import ladies_sampler +from .sage import sage_sampler +from .shadow import shadow_sampler + +__all__ = [ + 'ladies_sampler', + 'sage_sampler', + 'shadow_sampler', +] + +classes = __all__ diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py new file mode 100644 index 0000000..9cb252e --- /dev/null +++ b/cagnet/samplers/ladies.py @@ -0,0 +1,112 @@ +import math +import torch +import torch.distributed as dist +import torch_sparse +from collections import defaultdict +from cagnet.samplers.utils import * + +from sparse_coo_tensor_cpp import downsample_gpu, compute_darts_gpu, throw_darts_gpu, \ + compute_darts_select_gpu, throw_darts_select_gpu, \ + compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ + shift_rowselect_gpu, shift_colselect_gpu, \ + scatteri_add_gpu, rowselect_coo_gpu + + +def start_time(timer): + timer.record() + +def stop_time(start_timer, stop_timer, barrier=False): + stop_timer.record() + torch.cuda.synchronize() + time_taken = start_timer.elapsed_time(stop_timer) + if barrier: + dist.barrier() + return time_taken + +def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ + semibulk_size, replication, sa_masks, rank, size, \ + row_groups, col_groups, timing, replicate_graph): + + total_start_timer = torch.cuda.Event(enable_timing=True) + total_stop_timer = torch.cuda.Event(enable_timing=True) + + timing_dict = defaultdict(list) + + frontier_size = frontier_size[0] # one layer + current_frontier = torch.cuda.IntTensor(mb_count_total, batch_size + frontier_size) + + node_count = adj_matrix.size(0) + node_count_total = adj_matrix.size(1) + mb_count = batches.size(0) + + rank_c = rank // replication + rank_col = rank % replication + + n_darts_col = n_darts // replication + if rank_col == replication - 1: + n_darts_col = n_darts - (replication - 1) * n_darts_col + n_darts_col = n_darts + + # adj_matrices = [[None] * n_layers for x in range(mb_count)] # adj_matrices[i][j] -- mb i layer j + adj_matrices = [None] * n_layers # adj_matrices[i] -- bulk minibatch matrix for layer j + + start_time(total_start_timer) + for i in range(n_layers): + print(f"Sampling layer {i}", flush=True) + if i == 0: + nnz = batch_size + else: + nnz = current_frontier[0, :].size(0) + + batches = batches.to_sparse_csr() + # adj_matrix_sq = adj_matrix.square().to_sparse_csr() + adj_matrix_sq = adj_matrix.values().square() + adj_matrix_sq = torch.sparse_csr_tensor(adj_matrix.crow_indices(), adj_matrix.col_indices(), + adj_matrix_sq, adj_matrix.size()) + p = gen_prob_dist(batches, adj_matrix_sq, mb_count, node_count_total, replication, rank, size, \ + row_groups, col_groups, sa_masks, timing_dict, "ladies", timing, replicate_graph) + + next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ + row_groups, col_groups, timing_dict, "ladies") + # adj_matrix = adj_matrix.to_sparse_coo() + if p.layout == torch.sparse_csr: + p = p.to_sparse_coo() + batches = batches.to_sparse_coo() + + batches_select, next_frontier_select, adj_matrix_sample = \ + select(next_frontier, adj_matrix_sq, batches, sa_masks, nnz, \ + batch_size, frontier_size, mb_count, mb_count_total, node_count_total, \ + replication, rank, size, row_groups, col_groups, timing_dict, i, "ladies", + semibulk_size) + + # adj_matrices[i] = adj_matrix_sample + adj_matrices[i] = adj_matrix_sample.to_sparse_csr() + current_frontier = next_frontier + + + print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) + if timing: + for k, v in sorted(timing_dict.items()): + if (k.startswith("spgemm") and k != "spgemm-misc") or k == "probability-spgemm" or k == "row-select-spgemm" or k == "col-select-spgemm" or k == "sampling-iters": + # print(f"{k} times: {v}") + v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) + v_tens_recv = [] + for i in range(size): + v_tens_recv.append(torch.cuda.FloatTensor(1).fill_(0)) + dist.all_gather(v_tens_recv, v_tens) + + if rank == 0: + min_time = min(v_tens_recv).item() + max_time = max(v_tens_recv).item() + avg_time = sum(v_tens_recv).item() / size + med_time = sorted(v_tens_recv)[size // 2].item() + + print(f"{k} min: {min_time} max: {max_time} avg: {avg_time} med: {med_time}") + dist.barrier() + for k, v in timing_dict.items(): + if len(v) > 0: + avg_time = sum(v) / len(v) + else: + avg_time = -1.0 + print(f"{k} total_time: {sum(v)} avg_time {avg_time} len: {len(v)}") + return batches_select, next_frontier_select, adj_matrices diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py new file mode 100644 index 0000000..140510f --- /dev/null +++ b/cagnet/samplers/sage.py @@ -0,0 +1,250 @@ +import math +import torch +import torch.distributed as dist +import torch_sparse +import numpy as np +from collections import defaultdict +from cagnet.samplers.utils import * + +from sparse_coo_tensor_cpp import downsample_gpu, compute_darts_gpu, throw_darts_gpu, \ + compute_darts_select_gpu, throw_darts_select_gpu, \ + compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ + shift_rowselect_gpu, shift_colselect_gpu, \ + scatteri_add_gpu, rowselect_coo_gpu, \ + sparse_coo_tensor_gpu + + +timing = True +baseline_compare = True + +def start_time(timer, timing_arg=None): + if timing_arg is not None: + start_timing = timing_arg + else: + start_timing = timing + if start_timing: + timer.record() + +def stop_time(start_timer, stop_timer, barrier=False, timing_arg=None): + if timing_arg is not None: + start_timing = timing_arg + else: + start_timing = timing + if start_timing: + stop_timer.record() + torch.cuda.synchronize() + time_taken = start_timer.elapsed_time(stop_timer) + if barrier: + dist.barrier() + return time_taken + else: + return 0.0 + +def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total, n_layers, n_darts_list, \ + replication, sa_masks, rank, size, row_groups, col_groups, + timing_arg, baseline_arg, replicate_graph): + + global timing + global baseline_compare + + timing = timing_arg + baseline_compare = baseline_arg + + total_start_timer = torch.cuda.Event(enable_timing=True) + total_stop_timer = torch.cuda.Event(enable_timing=True) + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + timing_dict = defaultdict(list) + + start_time(start_timer) + node_count = adj_matrix.size(0) + node_count_total = adj_matrix.size(1) + mb_count = batches.size(0) + + rank_c = rank // replication + rank_col = rank % replication + + # adj_matrices = [[None] * n_layers for x in range(mb_count)] + # adj_matrices[i][j] -- mb i layer j + adj_matrices = [None] * n_layers # adj_matrices[i] -- bulk mb mtx for layer j + frontiers = [None] * (n_layers + 1) # frontiers[i] -- bulk mb frontiers for layer j + + gpu = torch.device(f"cuda:{torch.cuda.current_device()}") + + batches_expand_rows = torch.arange(mb_count * batch_size, dtype=torch.int32, device=gpu) + batches_expand_idxs = torch.stack((batches_expand_rows, batches._indices()[1, :])) + # batches_expand = torch.sparse_coo_tensor( + # batches_expand_idxs, + # batches._values(), + # size=(mb_count * batch_size, node_count_total)) + print(f"batches_expand_idxs.dtype: {batches_expand_idxs.dtype}", flush=True) + batches_expand = sparse_coo_tensor_gpu(batches_expand_idxs, batches._values(), + torch.Size([mb_count * batch_size, node_count_total])) + + if not replicate_graph: + batches_expand = batches_expand.to_sparse_csr() + timing_dict["sage-preamble"].append(stop_time(start_timer, stop_timer)) + + # adj_matrix = adj_matrix.to_sparse_csr() + current_frontier = batches_expand + + for i in range(n_layers): + start_time(total_start_timer) + # print(f"Sampling layer {i}", flush=True) + if i == 0: + nnz = batch_size + else: + # Restructure current_frontier to only have 1 nnz/row + # current_frontier_nnzmask = current_frontier._values().nonzero().squeeze() + # current_frontier_nnzcols = current_frontier._indices()[1, current_frontier_nnzmask] + # current_frontier_nnzrows = torch.arange(current_frontier_nnzcols.numel()).cuda() + start_time(start_timer) + current_frontier_nnzcols = current_frontier._indices()[1, :] + current_frontier_nnzrows = torch.arange(current_frontier._nnz()).cuda() + current_frontier_nnzinds = torch.stack((current_frontier_nnzrows, current_frontier_nnzcols)) + # current_frontier_nnzvals = current_frontier._values()[current_frontier_nnzmask].double() + #current_frontier_nnzvals = current_frontier._values().double() + current_frontier_nnzvals = current_frontier._values().float() + # current_frontier = torch.sparse_coo_tensor(current_frontier_nnzinds, current_frontier_nnzvals, + # size=torch.Size([current_frontier._nnz(), + # # size=torch.Size([current_frontier_nnzcols.numel(), + # current_frontier.size(1)])) + current_frontier = sparse_coo_tensor_gpu(current_frontier_nnzinds, current_frontier_nnzvals, + torch.Size([current_frontier._nnz(), current_frontier.size(1)])) + # nnz = current_frontier._nnz() // mb_count + # nnz = batch_size * (frontier_size ** i) + nnz = batch_size * int(np.prod(frontier_sizes[:i], dtype=int)) + if not replicate_graph: + current_frontier = current_frontier.to_sparse_csr() + timing_dict["sage-startiter"].append(stop_time(start_timer, stop_timer)) + + # Expand batches matrix + p = gen_prob_dist(current_frontier, adj_matrix, mb_count, node_count_total, + replication, rank, size, row_groups, col_groups, + sa_masks, timing_dict, "sage", + timing_arg, replicate_graph) + + # dist.barrier() + + # adj_matrix = adj_matrix.to_sparse_coo() + # start_time(start_timer) + # if p.layout == torch.sparse_csr: + # p = p.to_sparse_coo() + # timing_dict["sage-csr2coo"].append(stop_time(start_timer, stop_timer)) + + frontier_size = frontier_sizes[i] + n_darts = n_darts_list[i] + next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, + replication, rank, size, row_groups, col_groups, + timing_dict, "sage") + # dist.barrier() + + start_time(start_timer) + # add explicit 0's to next_frontier to fix the number of rows as bs * sn^l + next_frontier_nnz = next_frontier._values().nonzero().squeeze() + frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) + + frontier_nnz_sizes = torch.clamp(frontier_nnz_sizes, max=frontier_size) + next_frontier_rows = torch.repeat_interleave( + torch.arange(nnz * mb_count, device=gpu), + frontier_size) + nextf_cols_idxs = torch.arange(next_frontier_nnz.size(0), device=gpu) + frontier_remainder = frontier_size - frontier_nnz_sizes + ps_f_remain = torch.cumsum(frontier_remainder, dim=0).roll(1) + ps_f_remain[0] = 0 + nextf_cols_idxs += torch.repeat_interleave(ps_f_remain, frontier_nnz_sizes) + next_frontier_cols = torch.cuda.LongTensor(next_frontier_rows.size(0)).fill_(0) + next_frontier_cols.scatter_(0, nextf_cols_idxs, next_frontier._indices()[1,next_frontier_nnz]) + + next_frontier_idxs = torch.stack((next_frontier_rows, next_frontier_cols)) + next_frontier_values = torch.cuda.LongTensor(next_frontier_rows.size(0)).fill_(0) + next_frontier_values[nextf_cols_idxs] = 1 + + # Construct sampled adj matrix + # next_frontier = torch.sparse_coo_tensor(next_frontier_idxs, + # next_frontier_values, + # size=(nnz * mb_count, node_count_total)) + # # size=(batch_size * mb_count, node_count_total)) + next_frontier = sparse_coo_tensor_gpu(next_frontier_idxs, next_frontier_values, + torch.Size([nnz * mb_count, node_count_total])) + + next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) + # current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ + # current_frontier.values().bool()).view(current_frontier._nnz(), 1) + timing_dict["frontier-row-col-select"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + if not replicate_graph: + current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) + else: + current_frontier_select = current_frontier._indices()[1,:].view(current_frontier._nnz(), 1) + if i == 0: + # frontiers[0] = current_frontier_select.clone() + frontiers[0] = current_frontier_select + + batch_vals = torch.cuda.LongTensor(current_frontier_select.size()).fill_(1) + next_frontier_select_vals = next_frontier._values().view(-1) + + batch_rows = torch.arange(mb_count * nnz, device=gpu).view(mb_count * nnz, 1) + next_frontier_select_rows = next_frontier._indices()[0,:].view(-1) + + nnz_mask = next_frontier_select_vals.nonzero().squeeze() + adj_sample_rows = next_frontier_select_rows[nnz_mask] + adj_sample_cols = torch.arange(next_frontier_select.numel(), device=gpu) + adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * nnz) + adj_sample_cols = adj_sample_cols[nnz_mask] + adj_matrices_indices = torch.stack((adj_sample_rows, adj_sample_cols)) + # adj_matrices_values = torch.cuda.DoubleTensor(adj_sample_rows.size(0)).fill_(1) + adj_matrices_values = torch.cuda.FloatTensor(adj_sample_rows.size(0)).fill_(1) + + # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, + # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) + adj_matrix_sample = sparse_coo_tensor_gpu(adj_matrices_indices, adj_matrices_values, + torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) + # adj_matrices[i] = adj_matrix_sample + # frontiers[i + 1] = next_frontier_select.clone() + if i + 1 == n_layers: + frontiers[i + 1] = next_frontier_select + # else: + # del next_frontier_select + # del p + + # if i > 0: + # del current_frontier_select + current_frontier = next_frontier + timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + adj_matrices[i] = adj_matrix_sample.to_sparse_csr() + timing_dict["adj-coo2csr"].append(stop_time(start_timer, stop_timer)) + # del next_frontier_select + # del current_frontier_select + timing_dict["sage-samplingiter"].append(stop_time(total_start_timer, total_stop_timer)) + + if timing: + for k, v in sorted(timing_dict.items()): + if (k.startswith("spgemm") and k != "spgemm-misc") or k == "probability-spgemm" or k == "row-select-spgemm" or k == "col-select-spgemm" or k == "sampling-iters" or k == "frontier-row-col-select" or k == "adj-row-col-select" or k.startswith("sample") or k == "compute-p" or k == "sage-startiter" or k == "sage-csr2coo" or k == "sage-preamble" or k == "sage-samplingiter": + v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) + v_tens_recv = [] + for i in range(size): + v_tens_recv.append(torch.cuda.FloatTensor(1).fill_(0)) + dist.all_gather(v_tens_recv, v_tens) + + if rank == 0: + min_time = min(v_tens_recv).item() + max_time = max(v_tens_recv).item() + avg_time = sum(v_tens_recv).item() / size + med_time = sorted(v_tens_recv)[size // 2].item() + + print(f"{k} min: {min_time} max: {max_time} avg: {avg_time} med: {med_time}") + dist.barrier() + for k, v in timing_dict.items(): + if len(v) > 0: + avg_time = sum(v) / len(v) + else: + avg_time = -1.0 + print(f"{k} total_time: {sum(v)} avg_time {avg_time} len: {len(v)}") + # return current_frontier_select, next_frontier_select, adj_matrices + return frontiers, adj_matrices diff --git a/cagnet/samplers/shadow.py b/cagnet/samplers/shadow.py new file mode 100644 index 0000000..f544ad3 --- /dev/null +++ b/cagnet/samplers/shadow.py @@ -0,0 +1,244 @@ +import math +import torch +import torch.distributed as dist +import torch_sparse +import numpy as np +from collections import defaultdict +from cagnet.samplers.utils import * +from sparse_coo_tensor_cpp import * + +timing = True + +# def start_time(timer, timing_arg=None): +# if timing_arg is not None: +# start_timing = timing_arg +# else: +# start_timing = timing +# if start_timing: +# timer.record() +# +# def stop_time(start_timer, stop_timer, stream=None, barrier=False, timing_arg=None): +# if timing_arg is not None: +# start_timing = timing_arg +# else: +# start_timing = timing +# if start_timing: +# stop_timer.record() +# if stream is None: +# torch.cuda.synchronize() +# else: +# stream.synchronize() +# time_taken = start_timer.elapsed_time(stop_timer) +# if barrier: +# dist.barrier() +# return time_taken +# else: +# return 0.0 + +def shadow_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total, n_layers, n_darts_list, \ + replication, rank, size, row_groups, col_groups, + timing_arg, timing_dict, replicate_graph, epoch): + + global timing + + timing = timing_arg + + total_start_timer = torch.cuda.Event(enable_timing=True) + total_stop_timer = torch.cuda.Event(enable_timing=True) + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + timing_dict_inner = defaultdict(list) + + rank_c = rank // replication + rank_col = rank % replication + mb_count = batches.size(0) + node_count = adj_matrix.size(0) + node_count_total = adj_matrix.size(1) + gpu = torch.device(f"cuda:{torch.cuda.current_device()}") + + # adj_matrix = adj_matrix.t().to_sparse_csr() + + batches_expand_rows = torch.arange(mb_count * batch_size, dtype=torch.int32, device=gpu) + batches_expand_idxs = torch.stack((batches_expand_rows, batches._indices()[1, :])) + batches_expand = sparse_coo_tensor_gpu(batches_expand_idxs, batches._values(), + torch.Size([mb_count * batch_size, node_count_total])) + + + current_frontier = batches_expand + + # frontiers = torch.cuda.LongTensor(mb_count * batch_size, node_count_total).fill_(0) # TODO: make sparse + # frontiers_idxs = batches._indices()[1,:].unsqueeze(1) + # frontiers.scatter_(1, frontiers_idxs, 1) + frontiers_rows = torch.arange(mb_count * batch_size, dtype=torch.int64).to(gpu) + frontiers_cols = batches._indices()[1,:] + frontiers_idxs = torch.stack((frontiers_rows, frontiers_cols)) + frontiers_vals = torch.cuda.LongTensor(mb_count * batch_size).fill_(1) + frontiers = torch.sparse_coo_tensor(frontiers_idxs, frontiers_vals, + size=torch.Size([mb_count * batch_size, node_count_total])) + current_frontier_mbids = None + if epoch >= 1: + start_time(start_timer) + for i in range(n_layers): + neighbors = batch_size * int(np.prod(frontier_sizes[:(i + 1)], dtype=int)) + + p = gen_prob_dist(current_frontier, adj_matrix, mb_count, node_count_total, + replication, rank, size, row_groups, col_groups, + None, timing_dict_inner, "sage", timing_arg, replicate_graph) + + # frontier_size = frontier_sizes[i] + frontier_size = frontier_sizes[0] + n_darts = n_darts_list[i] + next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, + replication, rank, size, row_groups, col_groups, + timing_dict_inner, "sage") + + # collapse next_frontier to mb_count x node_count_total sparse matrix + # TODO: Change this when collapsed frontier is sparse + next_frontier_mask = next_frontier._values().nonzero().squeeze() + # collapsed_frontier_rows = next_frontier._indices()[0, next_frontier_mask] + + if i == 0: + collapsed_frontier_rows = next_frontier._indices()[0, next_frontier_mask] + else: + collapsed_frontier_rows = current_frontier_mbids[next_frontier._indices()[0, next_frontier_mask]] + + # collapsed_frontier_rows = collapsed_frontier_rows.div(row_div, rounding_mode="floor") + collapsed_frontier_cols = next_frontier._indices()[1, next_frontier_mask] + collapsed_frontier_idxs = torch.stack((collapsed_frontier_rows, collapsed_frontier_cols)) + if collapsed_frontier_rows.numel() == 1 and collapsed_frontier_cols.numel() == 1: + collapsed_frontier_idxs = collapsed_frontier_idxs.reshape(2, 1) + collapsed_frontier_vals = torch.cuda.LongTensor(collapsed_frontier_idxs.size(1)).fill_(1) + + collapsed_frontier = torch.sparse_coo_tensor(collapsed_frontier_idxs, + collapsed_frontier_vals, + torch.Size([mb_count * batch_size, node_count_total])) + collapsed_frontier = collapsed_frontier.coalesce() + # collapsed_frontier_dense = collapsed_frontier.to_dense() + # frontiers = frontiers + collapsed_frontier_dense + current_frontier_mbids = collapsed_frontier_idxs[0,:] + frontiers = frontiers + collapsed_frontier + frontiers = frontiers.coalesce() + + if i < n_layers - 1: + # expand_frontier_rows = torch.arange(collapsed_frontier._nnz()).to(gpu) + if collapsed_frontier_cols.dim() == 0: + collapsed_frontier_cols = collapsed_frontier_cols.unsqueeze(0) + expand_frontier_rows = torch.arange(collapsed_frontier_cols.size(0)).to(gpu) + expand_frontier_cols = collapsed_frontier_cols + expand_frontier_vals = collapsed_frontier_vals + expand_frontier_idxs = torch.stack((expand_frontier_rows, expand_frontier_cols)) + expand_frontier = torch.sparse_coo_tensor(expand_frontier_idxs, + expand_frontier_vals, + torch.Size([mb_count * neighbors, node_count_total])) + current_frontier = expand_frontier + if epoch >= 1: + timing_dict["shadow-sampling"].append(stop_time(start_timer, stop_timer)) + + if epoch >= 1: + start_time(start_timer) + # sampled_frontiers = frontiers.nonzero()[:,1] + sampled_frontiers = frontiers._indices()[1,:] + + batch_ids = torch.div(frontiers._indices()[0,:], batch_size, rounding_mode="floor").int() + batch_sizes = torch.histc(batch_ids, bins=mb_count, min=0, max=mb_count-1) + ps_batch_sizes = torch.cumsum(batch_sizes, 0).roll(1) + ps_batch_sizes[0] = 0 + + row_lengths = adj_matrix.crow_indices()[1:] - adj_matrix.crow_indices()[:-1] + rowselect_adj_crows = torch.cuda.LongTensor(sampled_frontiers.size(0) + 1).fill_(0) + rowselect_adj_crows[1:] = torch.cumsum(row_lengths[sampled_frontiers], dim=0) + rowselect_adj_crows[0] = 0 + rowselect_adj_nnz = rowselect_adj_crows[-1].item() + rowselect_adj_cols = torch.cuda.LongTensor(rowselect_adj_nnz).fill_(0) + rowselect_adj_vals = torch.cuda.LongTensor(rowselect_adj_nnz).fill_(0) + rowselect_csr_dupes_gpu(sampled_frontiers, adj_matrix.crow_indices().long(), adj_matrix.col_indices(), + adj_matrix.values(), rowselect_adj_crows, rowselect_adj_cols, rowselect_adj_vals, + sampled_frontiers.size(0), adj_matrix._nnz()) + + sampled_frontier_size = sampled_frontiers.size(0) + row_select_adj = torch.sparse_csr_tensor(rowselect_adj_crows, rowselect_adj_cols, rowselect_adj_vals, + torch.Size([sampled_frontier_size, node_count_total])) + + # sampled_frontiers_rowids = frontiers.nonzero()[:,0] + sampled_frontiers_rowids = frontiers._indices()[0,:] + sampled_frontiers_csr = frontiers.to_sparse_csr() + + colselect_idxs = torch.cuda.LongTensor(row_select_adj._nnz()).fill_(-1) + shadow_colselect_gpu(sampled_frontiers, sampled_frontiers_rowids, sampled_frontiers_csr.crow_indices(), + sampled_frontiers_csr.col_indices(), row_select_adj.crow_indices().long(), + row_select_adj.col_indices(), colselect_idxs, ps_batch_sizes, batch_size, + sampled_frontiers.size(0), row_select_adj._nnz()) + + colselect_mask = colselect_idxs > -1 + sampled_adjs_cols = row_select_adj.col_indices()[colselect_mask] + # for i in range(sampled_frontier_size): + # vtx = sampled_frontiers[i] + # frontiers_col_indices = frontiers._indices()[1, frontiers._indices()[0,:] == i] + # adj_col_indices = sampled_adjs_cols[row_select_adj.to_sparse_coo()._indices()[0, colselect_mask] == i] + # print(f"i: {i} vtx: {vtx} sampled_frontiers[i]: {frontiers_col_indices} sampled_adjs[i]: {adj_col_indices}", flush=True) + sampled_adjs_cols = colselect_idxs[colselect_mask] + sampled_adjs_vals = row_select_adj.values()[colselect_mask] + + sampled_adjs_rows = row_select_adj.to_sparse_coo()._indices()[0,:] + sampled_adjs_rows = sampled_adjs_rows[colselect_mask] + + sampled_adj_indices = torch.stack((sampled_adjs_rows, sampled_adjs_cols)) + sampled_frontier_size = sampled_frontiers.size(0) + sampled_adjs = torch.sparse_coo_tensor(sampled_adj_indices, sampled_adjs_vals, + torch.Size([sampled_frontier_size, sampled_frontier_size])) + + # sampled_adjs = sampled_adjs.t() + + sampled_frontiers_split = torch.split(sampled_frontiers, batch_sizes.tolist()) + if epoch >= 1: + timing_dict["shadow-selection"].append(stop_time(start_timer, stop_timer)) + + # sampled_adjs_indices_split = torch.split(sampled_adjs._indices(), batch_sizes.tolist(), dim=1) + # sampled_adjs_vals_split = torch.split(sampled_adjs._values(), batch_sizes.tolist()) + if epoch >= 1: + start_time(start_timer) + sampled_adjs_split = [] + for i in range(mb_count): + if i < mb_count - 1: + sampled_adjs_mask = (sampled_adjs._indices()[0,:] >= ps_batch_sizes[i]) & \ + (sampled_adjs._indices()[0,:] < ps_batch_sizes[i + 1]) + else: + sampled_adjs_mask = (sampled_adjs._indices()[0,:] >= ps_batch_sizes[i]) & \ + (sampled_adjs._indices()[0,:] < batch_sizes.sum().item()) + sampled_adjs_indices_split = sampled_adjs._indices()[:, sampled_adjs_mask] + # sampled_adjs_indices_split[0,:] -= sampled_adjs_indices_split[0,:].min().item() + sampled_adjs_indices_split[0,:] -= ps_batch_sizes[i] + sampled_adjs_values_split = sampled_adjs._values()[sampled_adjs_mask] + sampled_adjs_split.append(torch.sparse_coo_tensor(sampled_adjs_indices_split, + sampled_adjs_values_split, + torch.Size([batch_sizes[i], batch_sizes[i]])).t()) + if epoch >= 1: + timing_dict["shadow-extraction"].append(stop_time(start_timer, stop_timer)) + + # if timing: + # # for k, v in sorted(timing_dict.items()): + # # if (k.startswith("spgemm") and k != "spgemm-misc") or k == "probability-spgemm" or k == "row-select-spgemm" or k == "col-select-spgemm" or k == "sampling-iters" or k == "frontier-row-col-select" or k == "adj-row-col-select" or k.startswith("sample") or k == "compute-p" or k == "sage-startiter" or k == "sage-csr2coo" or k == "sage-preamble" or k == "sage-samplingiter": + # # v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) + # # v_tens_recv = [] + # # for i in range(size): + # # v_tens_recv.append(torch.cuda.FloatTensor(1).fill_(0)) + # # dist.all_gather(v_tens_recv, v_tens) + + # # if rank == 0: + # # min_time = min(v_tens_recv).item() + # # max_time = max(v_tens_recv).item() + # # avg_time = sum(v_tens_recv).item() / size + # # med_time = sorted(v_tens_recv)[size // 2].item() + + # # print(f"{k} min: {min_time} max: {max_time} avg: {avg_time} med: {med_time}") + # # dist.barrier() + # for k, v in timing_dict.items(): + # if len(v) > 0: + # avg_time = sum(v) / len(v) + # else: + # avg_time = -1.0 + # print(f"{k} total_time: {sum(v)} avg_time {avg_time} len: {len(v)}") + # # return sampled_frontiers, sampled_adjs + return sampled_frontiers_split, sampled_adjs_split diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py new file mode 100644 index 0000000..e5fbb3c --- /dev/null +++ b/cagnet/samplers/utils.py @@ -0,0 +1,3295 @@ +import math +import numpy as np +import torch +import torch.distributed as dist +import torch_geometric +from torch_geometric.data import Data, Dataset +import torch_sparse +from collections import defaultdict +from pathlib import Path +import os +from sklearn.metrics import auc, roc_curve +from tqdm import tqdm +from matplotlib import pyplot as plt +from atlasify import atlasify +import seaborn as sns +import yaml + +from sparse_coo_tensor_cpp import downsample_gpu, compute_darts_gpu, throw_darts_gpu, \ + compute_darts_select_gpu, throw_darts_select_gpu, \ + compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ + normalize_csr_gpu, shift_rowselect_gpu, shift_colselect_gpu, \ + scatteri_add_gpu, rowselect_coo_gpu, \ + rowselect_csr_gpu, sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu, \ + sparse_csr_tensor_gpu, nsparse_spgemm, rearrange_rows_gpu, \ + rearrangel_rows_gpu, reduce_sum_gpu, sum_csrd_gpu, sum_csri_gpu + + +timing = True + +def start_time(timer, timing_arg=None): + if timing_arg is not None: + start_timing = timing_arg + else: + start_timing = timing + if start_timing: + timer.record() + +def stop_time(start_timer, stop_timer, stream=None, barrier=False, timing_arg=None): + if timing_arg is not None: + start_timing = timing_arg + else: + start_timing = timing + if start_timing: + stop_timer.record() + if stream is None: + torch.cuda.synchronize() + else: + stream.synchronize() + time_taken = start_timer.elapsed_time(stop_timer) + if barrier: + dist.barrier(torch.distributed.group.WORLD) + return time_taken + else: + return 0.0 + +def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=False): + if timing_dict is not None: + timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) + +def csr_allreduce(mat, left, right, rank, name=None, alg=None, timing_dict=None): + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + while left < right: + start_time(start_timer) + group_size = right - left + 1 + mid_rank = (left + right) // 2 + + if rank <= mid_rank: + recv_rank = rank + (group_size // 2) + else: + recv_rank = rank - (group_size // 2) + + ops = [None, None] + nnz_send = torch.cuda.IntTensor(1).fill_(mat._nnz()) + nnz_recv = torch.cuda.IntTensor(1) + + ops[0] = dist.P2POp(dist.isend, nnz_send, recv_rank) + ops[1] = dist.P2POp(dist.irecv, nnz_recv, recv_rank) + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + + # torch.cuda.synchronize() + rows_recv = torch.cuda.IntTensor(mat.size(0) + 1) + cols_recv = torch.cuda.IntTensor(nnz_recv.item()) + if mat.dtype == torch.float32: + vals_recv = torch.cuda.FloatTensor(nnz_recv.item()) + elif mat.dtype == torch.float64: + vals_recv = torch.cuda.DoubleTensor(nnz_recv.item()) + + ops = [None] * 6 + ops[0] = dist.P2POp(dist.isend, mat.crow_indices().int(), recv_rank, tag=0) + ops[1] = dist.P2POp(dist.isend, mat.col_indices().int(), recv_rank, tag=1) + ops[2] = dist.P2POp(dist.isend, mat.values(), recv_rank, tag=2) + ops[3] = dist.P2POp(dist.irecv, rows_recv, recv_rank, tag=0) + ops[4] = dist.P2POp(dist.irecv, cols_recv, recv_rank, tag=1) + ops[5] = dist.P2POp(dist.irecv, vals_recv, recv_rank, tag=2) + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-comm-{name}") + + # torch.cuda.synchronize() + start_time(start_timer) + torch.cuda.synchronize() + mat_recv = torch.sparse_csr_tensor(rows_recv.long(), cols_recv.long(), vals_recv, size=mat.size()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-csrinst-{name}") + start_time(start_timer) + + if True or alg == "ladies": + mat = mat + mat_recv + else: + # Add mat + mat_recv + mat_row_lens = mat.crow_indices()[1:] - mat.crow_indices()[:-1] + mat_recv_row_lens = mat_recv.crow_indices()[1:] - mat_recv.crow_indices()[:-1] + + mat_sum_row_lens = mat_row_lens + mat_recv_row_lens + mat_sum_crows = torch.cuda.LongTensor(mat_sum_row_lens.size(0) + 1).fill_(0) + mat_sum_crows[1:] = torch.cumsum(mat_sum_row_lens, dim=0) + mat_sum_crows[0] = 0 + mat_sum_cols = torch.cuda.LongTensor(mat._nnz() + mat_recv._nnz()).fill_(0) + mat_row_lens_nnzrows = mat_row_lens.nonzero().squeeze() + mat_recv_row_lens_nnzrows = mat_recv_row_lens.nonzero().squeeze() + reduce_sum_gpu(mat_sum_crows, mat.crow_indices(), mat_recv.crow_indices(), + mat_sum_cols, mat.col_indices(), mat_recv.col_indices(), + mat_row_lens_nnzrows, mat_recv_row_lens_nnzrows) + + mat_sum_vals = torch.cuda.FloatTensor(mat_sum_cols.size(0)).fill_(1.0) + mat_sum_crows = mat_sum_crows.cpu() + mat_sum_cols = mat_sum_cols.cpu() + mat_sum_vals = mat_sum_vals.cpu() + mat = torch.sparse_csr_tensor(mat_sum_crows, mat_sum_cols, mat_sum_vals, mat.size()) + mat = mat.to(device) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-sum-{name}") + + if rank <= mid_rank: + right = mid_rank + else: + left = mid_rank + 1 + return mat + +def col_select15D(sample_mtx, col_select_mtx, bulk_items, node_count_total, bulk_size, semibulk_size, + replication, rank, size, row_groups, col_groups, name, alg, timing_dict=None): + + rank_c = rank // replication + rank_col = rank % replication + stages = (bulk_items // semibulk_size) // replication + + sample_mtxs = [None] * stages + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + start_inner_timer = torch.cuda.Event(enable_timing=True) + stop_inner_timer = torch.cuda.Event(enable_timing=True) + + if alg == "sage": + mask = torch.cuda.BoolTensor(sample_mtx.size(1)).fill_(0) + col_indices_mask = torch.cuda.LongTensor(sample_mtx.size(1)).fill_(0) + idx_range = torch.arange(sample_mtx.size(1)).cuda() + + start_time(start_timer) + for i in range(stages): + q = stages * rank_col + i + + # Extract semibulk_size batches from sample_mtx + start_time(start_inner_timer) + batch_start = q * bulk_size * semibulk_size + batch_stop = (q + 1) * bulk_size * semibulk_size + batch_crow_indices = sample_mtx.crow_indices()[batch_start:(batch_stop + 1)].clone() + batch_crow_indices -= batch_crow_indices[0].item() + start_nnz = sample_mtx.crow_indices()[batch_start] + stop_nnz = sample_mtx.crow_indices()[batch_stop] + + batch_cols = sample_mtx.col_indices()[start_nnz:stop_nnz] + batch_vals = sample_mtx.values()[start_nnz:stop_nnz] + + batch_mtx = torch.sparse_csr_tensor(batch_crow_indices, batch_cols, batch_vals, \ + torch.Size([bulk_size * semibulk_size, sample_mtx.size(1)])) + + batch_mtx = batch_mtx.to_sparse_coo() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-batch-{name}", barrier=True) + + start_time(start_inner_timer) + # Extract semibulk_size column extraction matrices + chunk_row_start = q * semibulk_size * node_count_total + chunk_row_stop = (q + 1) * semibulk_size * node_count_total + chunk_row_mask = (col_select_mtx._indices()[0, :] >= chunk_row_start) & \ + (col_select_mtx._indices()[0, :] < chunk_row_stop) + col_chunk_indices = col_select_mtx._indices()[:, chunk_row_mask] + col_chunk_indices[0,:] -= chunk_row_start + col_chunk_values = col_select_mtx._values()[chunk_row_mask] + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-colsel-{name}", barrier=True) + + start_time(start_inner_timer) + if alg == "ladies": + sample_chunk_indices, sample_chunk_values = torch_sparse.spspmm(batch_mtx._indices(), \ + batch_mtx._values().double(), col_chunk_indices, \ + col_chunk_values, batch_mtx.size(0), \ + semibulk_size * node_count_total, \ + # col_select_mtx.size(1)) + col_select_mtx.size(1), coalesced=True) + elif alg == "sage": + mask.fill_(0) + batch_mtx = batch_mtx.to_sparse_coo() + mask[col_chunk_indices[0,:]] = True + col_mask = torch.gather(mask, 0, batch_mtx._indices()[1,:]) + + sample_rows = batch_mtx._indices()[0, col_mask] + sample_cols = batch_mtx._indices()[0, col_mask] + col_indices_mask.fill_(0) + col_indices_mask[col_chunk_indices[0,:]] = idx_range[:col_chunk_indices.size(1)].long() + sample_cols_scaled = torch.gather(col_indices_mask, 0, sample_cols) + sample_chunk_values = batch_mtx._values()[col_mask] + sample_chunk_indices = torch.stack((sample_rows, sample_cols_scaled)) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-spgemm-{name}", barrier=True) + sample_mtxs[i] = torch.sparse_coo_tensor(sample_chunk_indices, sample_chunk_values, + size=torch.Size([batch_mtx.size(0), col_select_mtx.size(1)])) + + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-extract-iters-{name}", barrier=True) + + start_time(start_timer) + sample_mtx_int = torch.cat(sample_mtxs, dim=0) + row_start = (bulk_items // replication) * bulk_size * rank_col + sample_mtx_rows = sample_mtx_int._indices()[0,:] + row_start + sample_mtx_indices = torch.stack((sample_mtx_rows, sample_mtx_int._indices()[1,:])) + sample_mtx_int = torch.sparse_coo_tensor(sample_mtx_indices, sample_mtx_int._values(), \ + size=torch.Size([bulk_items * bulk_size, col_select_mtx.size(1)])).to_sparse_csr() + + + rank_row_start = rank_c * replication + rank_row_stop = (rank_c + 1) * replication - 1 + sample_mtx_int = sample_mtx_int.to_sparse_csr() + sample_mtx_red = csr_allreduce(sample_mtx_int, rank_row_start, rank_row_stop, rank) + sample_mtx_red = sample_mtx_red.to_sparse_coo() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-extract-reduce-{name}", barrier=True) + + return sample_mtx_red._indices(), sample_mtx_red._values() + +# def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replication, rank, size, \ +# row_groups, col_groups, name, timing_dict=None): +# +# rank_c = rank // replication +# rank_col = rank % replication +# stages = mb_count // replication +# +# sample_mtxs = [None] * stages +# +# start_timer = torch.cuda.Event(enable_timing=True) +# stop_timer = torch.cuda.Event(enable_timing=True) +# +# start_inner_timer = torch.cuda.Event(enable_timing=True) +# stop_inner_timer = torch.cuda.Event(enable_timing=True) +# +# mask = torch.cuda.BoolTensor(sample_mtx.size(1)).fill_(0) +# col_indices_mask = torch.cuda.LongTensor(sample_mtx.size(1)).fill_(0) +# idx_range = torch.arange(sample_mtx.size(1)).long().cuda() +# start_time(start_timer) +# for i in range(stages): +# q = stages * rank_col + i +# +# start_time(start_inner_timer) +# next_frontier_unique = next_frontier_select[q].unique() +# batch_crow_indices = sample_mtx.crow_indices()[(q * batch_size):((q + 1) * batch_size + 1)].clone() +# batch_crow_indices -= batch_crow_indices[0].item() +# start_nnz = sample_mtx.crow_indices()[q * batch_size] +# stop_nnz = sample_mtx.crow_indices()[(q + 1) * batch_size] +# +# batch_cols = sample_mtx.col_indices()[start_nnz:stop_nnz] +# batch_vals = sample_mtx.values()[start_nnz:stop_nnz] +# +# batch_mtx = torch.sparse_csr_tensor(batch_crow_indices, batch_cols, batch_vals, \ +# torch.Size([batch_size, sample_mtx.size(1)])) +# stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-batch-{name}", barrier=True) +# +# start_time(start_inner_timer) +# mask.fill_(0) +# batch_mtx = batch_mtx.to_sparse_coo() +# mask[next_frontier_unique] = True +# col_mask = torch.gather(mask, 0, batch_mtx._indices()[1,:]) +# stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-gather-{name}", barrier=True) +# +# start_time(start_inner_timer) +# sample_rows = batch_mtx._indices()[0, col_mask] +# sample_cols = batch_mtx._indices()[1, col_mask] +# col_indices_mask.fill_(0) +# col_indices_mask[next_frontier_unique] = idx_range[:next_frontier_unique.size(0)] +# sample_cols_scaled = torch.gather(col_indices_mask, 0, sample_cols) +# sample_vals = batch_mtx._values()[col_mask] +# sample_indices = torch.stack((sample_rows, sample_cols_scaled)) +# sample_mtxs[i] = torch.sparse_coo_tensor(sample_indices, sample_vals, +# size=torch.Size([batch_size, next_frontier_select.size(1)])) +# stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-cols-{name}", barrier=True) +# +# # start_time(start_inner_timer) +# # batch_mtx = batch_mtx.t() +# # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-transpbatch-{name}", barrier=True) +# +# # start_time(start_inner_timer) +# # col_mask = torch.cuda.BoolTensor(batch_mtx._nnz()).fill_(0) +# # rowselect_csr_gpu(next_frontier_unique, batch_mtx.crow_indices(), col_mask, \ +# # next_frontier_unique.size(0), batch_mtx._nnz()) +# +# # sample_row_lengths = batch_mtx.crow_indices()[1:] - batch_mtx.crow_indices()[:-1] +# # sample_row_lengths = sample_row_lengths[next_frontier_unique] +# # sample_crow_indices = torch.cuda.LongTensor(next_frontier_unique.size(0) + 1) +# # sample_crow_indices[0] = 0 +# # sample_crow_indices[1:] = torch.cumsum(sample_row_lengths, dim=0) +# # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-cols-{name}", barrier=True) +# # +# # start_time(start_inner_timer) +# # sample_cols = batch_mtx.col_indices()[col_mask] +# # sample_vals = batch_mtx.values()[col_mask] +# +# # sample_mtxs[i] = torch.sparse_csr_tensor(sample_crow_indices, sample_cols, sample_vals, \ +# # size=torch.Size([next_frontier_unique.size(0), batch_size])) +# # sample_mtxs[i] = sample_mtxs[i].t().to_sparse_coo() +# # sample_mtxs[i] = torch.sparse_coo_tensor(sample_mtxs[i]._indices(), sample_mtxs[i]._values(), +# # size=torch.Size([batch_size, next_frontier_select.size(1)])) +# # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-sample-{name}", barrier=True) +# +# stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-extract-iters-{name}", barrier=True) +# +# start_time(start_timer) +# sample_mtx = torch.cat(sample_mtxs, dim=0) +# row_start = (mb_count // replication) * batch_size * rank_col +# sample_mtx_rows = sample_mtx._indices()[0,:] + row_start +# sample_mtx_indices = torch.stack((sample_mtx_rows, sample_mtx._indices()[1,:])) +# sample_mtx = torch.sparse_coo_tensor(sample_mtx_indices, sample_mtx._values(), \ +# size=torch.Size([mb_count * batch_size, next_frontier_select.size(1)])).to_sparse_csr() +# +# +# rank_row_start = rank_c * replication +# rank_row_stop = (rank_c + 1) * replication - 1 +# sample_mtx = sample_mtx.to_sparse_csr() +# sample_mtx = csr_allreduce(sample_mtx, rank_row_start, rank_row_stop, rank) +# sample_mtx = sample_mtx.to_sparse_coo() +# stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-extract-reduce-{name}", barrier=True) +# +# return sample_mtx._indices(), sample_mtx._values() + +def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, name, timing_dict=None): + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + start_time(start_timer) + chunk_size = math.ceil(float(mata.size(1) / (size / replication))) + matc = torch.sparse_coo_tensor(size=(mata.size(0), matb.size(1))).cuda() + rank_c = rank // replication + rank_col = rank % replication + stages = size // (replication ** 2) + if rank_col == replication - 1: + stages = (size // replication) - (replication - 1) * stages + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-inst-{name}", barrier=True) + + for i in range(stages): + start_time(start_timer) + q = (rank_col * (size // (replication ** 2)) + i) * replication + rank_col + + matb_recv_nnz = torch.cuda.IntTensor([matb._nnz()]) + dist.broadcast(matb_recv_nnz, src=q, group=col_groups[rank_col]) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-bcast-nnz-{name}", barrier=True) + + start_time(start_timer) + if q == rank: + matb_recv_indices = matb._indices().clone() + matb_recv_values = matb._values().clone() + else: + matb_recv_indices = torch.cuda.LongTensor(2, matb_recv_nnz.item()) + matb_recv_values = torch.cuda.DoubleTensor(matb_recv_nnz.item()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-inst-recv-{name}", barrier=True) + + start_time(start_timer) + dist.broadcast(matb_recv_indices, src=q, group=col_groups[rank_col]) + dist.broadcast(matb_recv_values, src=q, group=col_groups[rank_col]) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-bcast-data-{name}", barrier=True) + + start_time(start_timer) + am_partid = rank_col * (size // replication ** 2) + i + chunk_col_start = am_partid * chunk_size + chunk_col_stop = min((am_partid + 1) * chunk_size, mata.size(1)) + chunk_col_size = chunk_col_stop - chunk_col_start + chunk_col_mask = (mata._indices()[1, :] >= chunk_col_start) & (mata._indices()[1, :] < chunk_col_stop) + + mata_chunk_indices = mata._indices()[:, chunk_col_mask] + mata_chunk_indices[1,:] -= chunk_col_start + mata_chunk_values = mata._values()[chunk_col_mask] + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}", barrier=True) + + start_time(start_timer) + matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ + mata_chunk_values.double(), matb_recv_indices, \ + matb_recv_values, mata.size(0), \ + chunk_col_size, matb.size(1), coalesced=True) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) + + start_time(start_timer) + # matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + matc_chunk = sparse_coo_tensor_gpu(matc_chunk_indices, matc_chunk_values, matc.size()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}", barrier=True) + + start_time(start_timer) + matc_chunk = matc_chunk.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}", barrier=True) + + start_time(start_timer) + matc += matc_chunk + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}", barrier=True) + + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}", barrier=True) + + # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) + # Implement sparse allreduce w/ all_gather and padding + start_time(start_timer) + matc_nnz = torch.cuda.IntTensor(1).fill_(matc._nnz()) + dist.all_reduce(matc_nnz, dist.ReduceOp.MAX, row_groups[rank_c]) + matc_nnz = matc_nnz.item() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-nnz-{name}", barrier=True) + + start_time(start_timer) + matc_recv_indices = [] + matc_recv_values = [] + for i in range(replication): + matc_recv_indices.append(torch.cuda.LongTensor(2, matc_nnz).fill_(0)) + matc_recv_values.append(torch.cuda.DoubleTensor(matc_nnz).fill_(0.0)) + + matc_send_indices = torch.cat((matc._indices(), torch.cuda.LongTensor(2, matc_nnz - matc._nnz()).fill_(0)), 1) + matc_send_values = torch.cat((matc._values(), torch.cuda.DoubleTensor(matc_nnz - matc._nnz()).fill_(0.0))) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-padding-{name}", barrier=True) + + start_time(start_timer) + dist.all_gather(matc_recv_indices, matc_send_indices, row_groups[rank_c]) + dist.all_gather(matc_recv_values, matc_send_values, row_groups[rank_c]) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-allgather-{name}", barrier=True) + + start_time(start_timer) + matc_recv = [] + for i in range(replication): + # matc_recv.append(torch.sparse_coo_tensor(matc_recv_indices[i], matc_recv_values[i], matc.size())) + matc_recv.append(sparse_coo_tensor_gpu(matc_recv_indices[i], matc_recv_values[i], matc.size())) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}", barrier=True) + + start_time(start_timer) + matc_recv = torch.stack(matc_recv) + matc = torch.sparse.sum(matc_recv, dim=0) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}", barrier=True) + + start_time(start_timer) + nnz_mask = matc._values() > 0 + matc_nnz_indices = matc._indices()[:, nnz_mask] + matc_nnz_values = matc._values()[nnz_mask] + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-unpad-{name}", barrier=True) + + return matc_nnz_indices, matc_nnz_values + +def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ + name, nnz_row_masks, timing_dict=None, alg=None): + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + start_inner_timer = torch.cuda.Event(enable_timing=True) + stop_inner_timer = torch.cuda.Event(enable_timing=True) + + start_time(start_timer) + # chunk_size = math.ceil(float(mata.size(1) / (size / replication))) + chunk_size = mata.size(1) // (size // replication) + if True or mata.layout == torch.sparse_coo or (name == "prob" and alg == "sage"): + matc = torch.sparse_coo_tensor(size=(mata.size(0), matb.size(1))).cuda() + elif mata.layout == torch.sparse_csr: + matc_crow_indices = torch.cuda.LongTensor(mata.size(0) + 1).fill_(0) + matc = torch.sparse_csr_tensor(matc_crow_indices, \ + torch.cuda.LongTensor(0), \ + torch.cuda.FloatTensor(0), + size=(mata.size(0), matb.size(1))) + rank_c = rank // replication + rank_col = rank % replication + stages = size // (replication ** 2) + if rank_col == replication - 1: + stages = (size // replication) - (replication - 1) * stages + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-inst-{name}", barrier=True) + + for i in range(stages): + # start_time(start_timer) + q = (rank_col * (size // (replication ** 2)) + i) * replication + rank_col + + # Extract chunk of mata for local SpGEMM + start_time(start_timer) + am_partid = rank_col * (size // replication ** 2) + i + chunk_col_start = am_partid * chunk_size + chunk_col_stop = min((am_partid + 1) * chunk_size, mata.size(1)) + chunk_col_size = chunk_col_stop - chunk_col_start + if mata.layout == torch.sparse_coo: + chunk_col_mask = (mata._indices()[1, :] >= chunk_col_start) & \ + (mata._indices()[1, :] < chunk_col_stop) + mata_chunk_indices = mata._indices()[:, chunk_col_mask] + mata_chunk_indices[1,:] -= chunk_col_start + mata_chunk_values = mata._values()[chunk_col_mask] + elif mata.layout == torch.sparse_csr: + # Column selection + chunk_col_mask = (mata.col_indices() >= chunk_col_start) & \ + (mata.col_indices() < chunk_col_stop) + mata_chunk_cols = mata.col_indices()[chunk_col_mask] + mata_chunk_cols -= chunk_col_start + mata_chunk_values = mata.values()[chunk_col_mask] + + # # torch 1.13 supports offsets-based reduction + # mata_chunk_rowcount = torch.segment_reduce(chunk_col_mask.float(), \ + # "sum", \ + # offsets=mata.crow_indices()) + + lengths = mata.crow_indices()[1:] - mata.crow_indices()[:-1] + mata_chunk_rowcount = torch.segment_reduce(chunk_col_mask.float(), \ + "sum", \ + lengths=lengths) + mata_chunk_crows = torch.cuda.LongTensor(mata.size(0) + 1) + mata_chunk_crows[1:] = torch.cumsum(mata_chunk_rowcount.long(), dim=0) + mata_chunk_crows[0] = 0 + + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}", barrier=True) + + # Determine number of nonzero columns in chunk of mata + start_time(start_timer) + if mata.layout == torch.sparse_coo: + nnz_cols = torch.unique(mata_chunk_indices[1,:]) + elif mata.layout == torch.sparse_csr: + nnz_cols = torch.unique(mata_chunk_cols) + nnz_cols_count = torch.cuda.IntTensor(1).fill_(nnz_cols.size(0)) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-unique-{name}", barrier=True) + # timing_dict[f"spgemm-nnzcount-{name}"].append(nnz_cols.size(0)) + + # Gather nnz column counts on rank q + start_time(start_timer) + # print(f"nnz_cols_count: {nnz_cols_count} q: {q} rank_col: {rank_col}", flush=True) + torch.cuda.synchronize() + if rank == q: + nnz_cols_count_list = [] + for j in range(size // replication): + nnz_cols_count_list.append(torch.cuda.IntTensor(1).fill_(0)) + recv_rank = rank_col + j * replication + dist.gather(nnz_cols_count, nnz_cols_count_list, dst=q, group=col_groups[rank_col]) + else: + dist.gather(nnz_cols_count, dst=q, group=col_groups[rank_col]) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-gather-nnzcounts-{name}", barrier=True) + + # Rank q allocates recv buffer for nnz col ids + start_time(start_timer) + if rank == q: + nnz_col_ids = [] + for nnz_count in nnz_cols_count_list: + nnz_col_ids.append(torch.cuda.LongTensor(nnz_count.item()).fill_(0)) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-alloc-nnzbuff-{name}", barrier=True) + + # isend/irecv nnz col ids to rank q + start_time(start_timer) + if rank == q: + recv_objs = [] + for j in range(size // replication): + recv_rank = rank_col + j * replication + if recv_rank != q: + recv_objs.append(dist.recv(nnz_col_ids[j], src=recv_rank)) + else: + nnz_col_ids[j] = nnz_cols + else: + dist.send(nnz_cols, dst=q) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-colids-{name}", barrier=True) + + # Send selected rows count + start_time(start_timer) + nnz_count = 0 + if rank == q: + nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) + if matb.layout == torch.sparse_coo: + start_time(start_inner_timer) + # print(f"nnz_col_ids.size: {nnz_col_ids.size()}", flush=True) + # print(f"nnz_col_ids: {nnz_col_ids}", flush=True) + rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:].long(), nnz_row_masks, nnz_cols_count_list.int(), \ + matb._indices().size(1), size // replication) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") + recv_proc_count = ((size // replication) - 1) + send_ops = [None] * 4 * recv_proc_count + send_idx = 0 + + # start_time(start_inner_timer) + for j in range(size // replication): + recv_rank = rank_col + j * replication + # nnz_row_mask = nnz_row_masks[(j * matb._nnz()):((j + 1) * matb._nnz())] + nnz_row_mask = torch.cuda.BoolTensor(matb._nnz()).fill_(False) + + if matb.layout == torch.sparse_csr: + start_time(start_inner_timer) + rowselect_csr_gpu(nnz_col_ids[j], matb.crow_indices().long(), nnz_row_mask, \ + nnz_col_ids[j].size(0), matb._nnz()) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") + + start_time(start_inner_timer) + row_lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] + if nnz_col_ids[j].size(0) > 0: + matb_send_lengths = row_lengths[nnz_col_ids[j]] + else: + matb_send_lengths = torch.cuda.LongTensor(0) + + matb_send_cols = matb.col_indices()[nnz_row_mask] + matb_send_values = matb.values()[nnz_row_mask] + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-misc-{name}") + + start_time(start_inner_timer) + if recv_rank != q: + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_cols.size(0)) + nnz_count += 2 * matb_send_cols.size(0) + nnz_col_ids[j].size(0) + # dist.send(selected_rows_count, tag=0, dst=recv_rank) + # # dist.send(matb_send_crows, tag=1, dst=recv_rank) + # dist.send(matb_send_lengths, tag=1, dst=recv_rank) + # dist.send(matb_send_cols, tag=2, dst=recv_rank) + # dist.send(matb_send_values, tag=3, dst=recv_rank) + + # print(f"send selected_rows_count: {selected_rows_count}", flush=True) + send_ops[send_idx] = \ + dist.P2POp(dist.isend, selected_rows_count, recv_rank, tag=0) + send_ops[send_idx + 1] = \ + dist.P2POp(dist.isend, matb_send_lengths.int(), recv_rank, tag=1) + send_ops[send_idx + 2] = \ + dist.P2POp(dist.isend, matb_send_cols.int(), recv_rank, tag=2) + send_ops[send_idx + 3] = \ + dist.P2POp(dist.isend, matb_send_values.float(), recv_rank, tag=3) + send_idx += 4 + else: + # matb_recv_crows = matb_send_crows.clone() + matb_recv_cols = matb_send_cols.clone() + matb_recv_values = matb_send_values.clone() + if not (name == "prob" and alg == "sage"): + # matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) + matb_rows_sum = torch.cuda.IntTensor(matb.crow_indices().size(0) - 1).fill_(0) + if nnz_row_mask.any(): + matb_rows_sum[nnz_col_ids[j]] = matb_send_lengths + # matb_recv_crows = torch.cuda.LongTensor(matb.crow_indices().size(0)) + # matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dim=0) + matb_recv_crows = torch.cuda.IntTensor(matb.crow_indices().size(0)) + matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dtype=torch.int32, dim=0) + matb_recv_crows[0] = 0 + else: + # matb_recv_rows = torch.repeat_interleave(nnz_col_ids[j], matb_send_lengths) + # matb_recv_indices = torch.stack((matb_recv_rows, matb_recv_cols)) + # matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) + matb_rows_sum = torch.cuda.IntTensor(matb.crow_indices().size(0) - 1).fill_(0) + if nnz_row_mask.any(): + matb_rows_sum[nnz_col_ids[j]] = matb_send_lengths + # matb_recv_crows = torch.cuda.LongTensor(matb.crow_indices().size(0)) + matb_recv_crows = torch.cuda.IntTensor(matb.crow_indices().size(0)) + matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dtype=torch.int32, dim=0) + matb_recv_crows[0] = 0 + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-calls-{name}") + elif matb.layout == torch.sparse_coo: + # rowselect_coo_gpu(nnz_col_ids[j], matb._indices()[0,:], nnz_row_mask, \ + # nnz_col_ids[j].size(0), matb._indices().size(1)) + + start_time(start_inner_timer) + matb_send_indices = matb._indices()[:, nnz_row_mask] + matb_send_values = matb._values()[nnz_row_mask] + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-misc-{name}") + + start_time(start_inner_timer) + if recv_rank != q: + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_indices.size(1)) + nnz_count += 3 * matb_send_indices.size(1) + dist.send(selected_rows_count, tag=0, dst=recv_rank) + dist.send(matb_send_indices, tag=1, dst=recv_rank) + dist.send(matb_send_values, tag=2, dst=recv_rank) + + # send_ops[send_idx] = dist.P2POp(dist.isend, selected_rows_count, recv_rank, tag=0) + # send_ops[send_idx + 1] = dist.P2POp(dist.isend, matb_send_indices, recv_rank, tag=1) + # send_ops[send_idx + 2] = dist.P2POp(dist.isend, matb_send_values, recv_rank, tag=2) + # send_idx += 3 + else: + # matb_select_recv = matb_send.clone() + matb_recv_indices = matb_send_indices.clone() + matb_recv_values = matb_send_values.clone() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-calls-{name}") + + # if len(send_ops) > 0: + if len(send_ops) > 0 and matb.layout == torch.sparse_csr: + reqs = dist.batch_isend_irecv(send_ops) + for req in reqs: + req.wait() + else: + if f"spgemm-send-rowsel-{name}" not in timing_dict: + timing_dict[f"spgemm-send-rowsel-{name}"] = [] + + if f"spgemm-send-misc-{name}" not in timing_dict: + timing_dict[f"spgemm-send-misc-{name}"] = [] + + if f"spgemm-send-calls-{name}" not in timing_dict: + timing_dict[f"spgemm-send-calls-{name}"] = [] + + start_time(start_inner_timer) + selected_rows_count_recv = torch.cuda.IntTensor(1).fill_(-1) + recv_ops = [dist.P2POp(dist.irecv, selected_rows_count_recv, q, tag=0)] + reqs = dist.batch_isend_irecv(recv_ops) + for req in reqs: + req.wait() + torch.cuda.synchronize() + # print(f"recv selected_rows_count: {selected_rows_count_recv}", flush=True) + # dist.recv(selected_rows_count_recv, tag=0, src=q) + + if matb.layout == torch.sparse_csr: + # matb_recv_crows = torch.cuda.LongTensor(chunk_col_size + 1) + # dist.recv(matb_recv_crows, tag=1, src=q) + + # matb_recv_lengths = torch.cuda.LongTensor(nnz_cols.size(0)) + recv_ops = [] + matb_recv_lengths = torch.cuda.IntTensor(nnz_cols.size(0)) + # dist.recv(matb_recv_lengths, tag=1, src=q) + # print(f"recv matb_recv_lengths: {matb_recv_lengths}", flush=True) + recv_ops.append(dist.P2POp(dist.irecv, matb_recv_lengths, q, tag=1)) + # matb_recv_lengths = matb_recv_lengths.long() + + # matb_recv_cols = torch.cuda.LongTensor(selected_rows_count_recv.item()) + matb_recv_cols = torch.cuda.IntTensor(selected_rows_count_recv.item()) + # dist.recv(matb_recv_cols, tag=2, src=q) + recv_ops.append(dist.P2POp(dist.irecv, matb_recv_cols, q, tag=2)) + # matb_recv_cols = matb_recv_cols.long() + + # matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) + matb_recv_values = torch.cuda.FloatTensor(selected_rows_count_recv.item()) + # dist.recv(matb_recv_values, tag=3, src=q) + recv_ops.append(dist.P2POp(dist.irecv, matb_recv_values, q, tag=3)) + + reqs = dist.batch_isend_irecv(recv_ops) + for req in reqs: + req.wait() + torch.cuda.synchronize() + + # the spgemm is better with COO for SAGE prob + if not (name == "prob" and alg == "sage"): + matb_recv_lengths = matb_recv_lengths.long() + matb_recv_cols = matb_recv_cols.long() + matb_rows_sum = torch.cuda.LongTensor(chunk_col_size).fill_(0) + if matb_recv_lengths.size(0) > 0: + matb_rows_sum[nnz_cols] = matb_recv_lengths + matb_recv_crows = torch.cuda.LongTensor(chunk_col_size + 1) + matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dtype=torch.int32, dim=0) + matb_recv_crows[0] = 0 + else: + # matb_recv_rows = torch.repeat_interleave(nnz_cols, matb_recv_lengths) + # matb_recv_indices = torch.stack((matb_recv_rows, matb_recv_cols)) + # matb_rows_sum = torch.cuda.LongTensor(chunk_col_size).fill_(0) + matb_rows_sum = torch.cuda.IntTensor(chunk_col_size).fill_(0) + if matb_recv_lengths.size(0) > 0: + matb_rows_sum[nnz_cols] = matb_recv_lengths + # matb_recv_crows = torch.cuda.LongTensor(chunk_col_size + 1) + matb_recv_crows = torch.cuda.IntTensor(chunk_col_size + 1) + # matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dim=0) + matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dtype=torch.int32, dim=0) + matb_recv_crows[0] = 0 + + nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_lengths.size(0) + + elif matb.layout == torch.sparse_coo: + + # matb_indices_recv = matb_recv_buff[:(3 * selected_rows_count_recv.item())].view(3, -1) + matb_recv_indices = torch.cuda.LongTensor(2, selected_rows_count_recv.item()) + dist.recv(matb_recv_indices, tag=1, src=q) + + matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) + dist.recv(matb_recv_values, tag=2, src=q) + + # recv_ops = [dist.P2POp(dist.irecv, matb_recv_indices, q, tag=1), \ + # dist.P2POp(dist.irecv, matb_recv_values, q, tag=2)] + + # reqs = dist.batch_isend_irecv(recv_ops) + # for req in reqs: + # req.wait() + nnz_count += 3 * matb_recv_indices.size(1) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-calls-{name}") + + torch.cuda.synchronize() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowdata-{name}", barrier=True) + + if f"spgemm-recv-calls-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-calls-{name}"] = [] + + timing_dict[f"spgemm-send-rownnz-{name}"].append(nnz_count) + + # start_time(start_timer) + # # print(f"matb_select_recv.dtype: {matb_select_recv.dtype}") + # # matb_recv_indices = matb_select_recv[:2, :].long() + # # # matb_recv_values = matb_select_recv[2, :].double() + # # matb_recv_values = matb_select_recv[2, :] + + # # matb_recv_indices = matb_select_recv[:2, :].long() + # # matb_recv_values = matb_select_recv[2, :].double() + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-castrecv-{name}", barrier=True) + + start_time(start_timer) + if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: + # Can skip local spgemm for row-selection for any alg and for SAGE's probability spgemm + # if name == "prob" and alg == "sage": + # matc_chunk_crows = torch.cuda.LongTensor(mata.size(0)).fill_(0) + # matc_chunk = torch.sparse_csr_tensor(matb_recv_crows, matb_recv_cols, matb_recv_values.float(), + # size=(chunk_col_size, matb.size(1))) + # matc += matc_chunk + if mata.layout == torch.sparse_csr and (name == "prob" and alg == "sage") or (name == "rowsel"): + start_time(start_inner_timer) + # mata_recv = torch.sparse_csr_tensor(mata_chunk_crows, mata_chunk_cols, mata_chunk_values, + # size=(mata.size(0), chunk_col_size)) + mata_recv_row_lens = mata_chunk_crows[1:] - mata_chunk_crows[:-1] + mata_recv_rows = torch.repeat_interleave( + torch.arange(0, mata.size(0), device=mata.device), + mata_recv_row_lens) + # matb_recv = sparse_coo_tensor_gpu(matb_recv_indices, matb_recv_values.float(), + # torch.Size([chunk_col_size, matb.size(1)])) + # matb_recv = matb_recv.to_sparse_csr() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-csrinst-{name}") + + + + start_time(start_inner_timer) + # mata_recv = mata_recv.to_sparse_coo() + # matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_recv._indices(), \ + # mata_recv._values().float(), matb_recv_indices, \ + # matb_recv_values.float(), mata.size(0), \ + # chunk_col_size, matb.size(1), coalesced=True) + # print(f"mata_nnz: {mata_chunk_values.size()} matb_nnz: {matb_recv._nnz()} dims: {mata.size(0)} {chunk_col_size} {matb_recv.size(1)}", flush=True) + + # matb_recv_row_lens = matb_recv.crow_indices()[1:] - matb_recv.crow_indices()[:-1] + matb_recv_row_lens = matb_recv_crows[1:] - matb_recv_crows[:-1] + # matc_chunk_row_lens = torch.cuda.LongTensor(mata.size(0) + 1).fill_(0) + matc_chunk_row_lens = torch.cuda.IntTensor(mata.size(0) + 1).fill_(0) + matc_chunk_row_lens[mata_recv_rows] = matb_recv_row_lens[mata_chunk_cols].int() + matc_chunk_crows = torch.cumsum(matc_chunk_row_lens, 0).roll(1) + matc_chunk_crows[0] = 0 + + # matc_chunk_cols = torch.cuda.LongTensor(matc_chunk_crows[-1].item()).fill_(0) + matc_chunk_cols = torch.cuda.IntTensor(matc_chunk_crows[-1].item()).fill_(0) + rearrange_rows_gpu(mata_recv_rows, mata_chunk_cols, matc_chunk_crows, + # rearrangel_rows_gpu(mata_recv_rows, mata_chunk_cols, matc_chunk_crows, + matb_recv_crows.int(), matb_recv_cols.int(), matc_chunk_cols) + # matb_recv.crow_indices(), matb_recv.col_indices(), matc_chunk_cols) + matc_chunk_values = torch.cuda.FloatTensor(matc_chunk_crows[-1].item()).fill_(1.0) + + # matc_chunk = nsparse_spgemm(mata_chunk_crows.int(), \ + # mata_chunk_cols.int(), \ + # mata_chunk_values.float(), \ + # matb_recv.crow_indices().int(), \ + # matb_recv.col_indices().int(), \ + # matb_recv.values().float(), \ + # mata.size(0), chunk_col_size, matb_recv.size(1)) + # matc_chunk_crows = matc_chunk[0] + # matc_chunk_cols = matc_chunk[1] + # matc_chunk_values = matc_chunk[2] + + # mata_recv = mata_recv.to_sparse_coo() + # matb_recv = matb_recv.to_sparse_coo() + # matc_chunk = torch.sparse.mm(mata_recv.float(), matb_recv.float()) + # matc_chunk_crows = matc_chunk.crow_indices() + # matc_chunk_cols = matc_chunk.col_indices() + # matc_chunk_values = matc_chunk.values() + matc_row_lens = matc_chunk_crows[1:] - matc_chunk_crows[:-1] + matc_chunk_rows = torch.repeat_interleave( + torch.arange(0, mata.size(0), device=mata.device), + matc_row_lens) + matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + + matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + matc_chunk = matc_chunk.coalesce() + matc += matc_chunk + torch.cuda.synchronize() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-spspmm-{name}") + # matc = matc.to_sparse_csr() + elif mata.layout == torch.sparse_csr: + start_time(start_inner_timer) + mata_recv = torch.sparse_csr_tensor(mata_chunk_crows, mata_chunk_cols, mata_chunk_values.float(), + size=(mata.size(0), chunk_col_size)) + matb_recv = torch.sparse_csr_tensor(matb_recv_crows, matb_recv_cols, matb_recv_values.float(), + size=(chunk_col_size, matb.size(1))) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-csrinst-{name}") + + + start_time(start_inner_timer) + mata_recv = mata_recv.to_sparse_coo() + matb_recv = matb_recv.to_sparse_coo() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-csr2coo-{name}") + + # matc += torch.mm(mata_recv, matb_recv) + start_time(start_inner_timer) + matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_recv._indices(), \ + mata_recv._values(), matb_recv._indices(), \ + matb_recv._values(), mata.size(0), \ + chunk_col_size, matb.size(1), coalesced=True) + + # mata_row_lens = mata_chunk_crows[1:] - mata_chunk_crows[:-1] + # torch.cuda.synchronize() + # print(f"before nsparse", flush=True) + # matc_chunk = nsparse_spgemm(mata_chunk_crows.int(), \ + # mata_chunk_cols.int(), \ + # mata_chunk_values.float(), \ + # matb_recv_crows.int(), \ + # matb_recv_cols.int(), \ + # matb_recv_values.float(), \ + # mata.size(0), chunk_col_size, matb.size(1)) + # print(f"after nsparse", flush=True) + # matc_chunk_crows = matc_chunk[0] + # matc_chunk_cols = matc_chunk[1] + # matc_chunk_values = matc_chunk[2] + + # matc_row_lens = matc_chunk_crows[1:] - matc_chunk_crows[:-1] + # matc_chunk_rows = torch.repeat_interleave( + # torch.arange(0, mata.size(0), device=mata.device), + # matc_row_lens) + # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + + matc_chunk = sparse_coo_tensor_gpu(matc_chunk_indices, matc_chunk_values, + torch.Size([matc.size(0), matc.size(1)])) + matc_chunk = matc_chunk.coalesce() + + matc += matc_chunk + matc_chunk = matc_chunk.to_sparse_csr() + # matc += torch.mm(mata_recv, matb_recv) + # matc_offsets, matc_cols, matc_values = spgemm_gpu(mata_recv.crow_indices().int(), mata_recv.col_indices().int(), mata_recv.values(), matb_recv.crow_indices().int(), matb_recv.col_indices().int(), matb_recv.values(), matc.crow_indices().int(), matc.col_indices().int(), matc.values(), int(mata_recv.size(0)), int(mata_recv.size(1)), int(matb_recv.size(1))) + # matc = torch.sparse_csr_tensor(matc_offsets, matc_cols, matc_values, size=matc.size()) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-spspmm-{name}") + # matc = matc.to_sparse_csr() + elif mata.layout == torch.sparse_coo: + matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ + mata_chunk_values, matb_recv_indices, \ + matb_recv_values, mata.size(0), \ + chunk_col_size, matb.size(1), coalesced=True) + matc_chunk = sparse_coo_tensor_gpu(matc_chunk_indices, matc_chunk_values, + torch.Size([matc.size(0), matc.size(1)])) + # matc_chunk = matc_chunk.coalesce() + matc += matc_chunk + if f"spgemm-loc-csrinst-{name}" not in timing_dict: + timing_dict[f"spgemm-loc-csrinst-{name}"] = [] + + if f"spgemm-loc-csr2coo-{name}" not in timing_dict: + timing_dict[f"spgemm-loc-csr2coo-{name}"] = [] + + if f"spgemm-loc-spspmm-{name}" not in timing_dict: + timing_dict[f"spgemm-loc-spspmm-{name}"] = [] + else: + if f"spgemm-loc-csrinst-{name}" not in timing_dict: + timing_dict[f"spgemm-loc-csrinst-{name}"] = [] + + # if f"spgemm-loc-csr2coo-{name}" not in timing_dict: + # timing_dict[f"spgemm-loc-csr2coo-{name}"] = [] + + if f"spgemm-loc-spspmm-{name}" not in timing_dict: + timing_dict[f"spgemm-loc-spspmm-{name}"] = [] + + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) + + # print(f"send-rowdata: {timing_dict['spgemm-send-rowdata-prob']}") + # print(f"send-rownnz: {timing_dict['spgemm-send-rownnz-prob']}") + # print(f"nnzcount: {timing_dict['spgemm-nnzcount-prob']}") + + rank_row_start = rank_c * replication + rank_row_stop = (rank_c + 1) * replication - 1 + + start_time(start_timer) + matc = matc.to_sparse_csr() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coo2csr{name}", barrier=True) + + start_time(start_timer) + # matc = csr_allreduce(matc, rank_row_start, rank_row_stop, rank) + matc = csr_allreduce(matc, rank_row_start, rank_row_stop, rank, "prob", alg, timing_dict) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) + + start_time(start_timer) + matc = matc.to_sparse_coo() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-csr2coo{name}", barrier=True) + return matc._indices(), matc._values() + + # # Old sparse reduction + # start_time(start_timer) + # if matc.layout == torch.sparse_csr: + # matc = matc.to_sparse_coo() + # matc = matc.coalesce() + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}", barrier=True) + + # # Implement sparse allreduce w/ all_gather and padding + # start_time(start_timer) + # matc_nnz = torch.cuda.IntTensor(1).fill_(matc._nnz()) + # dist.all_reduce(matc_nnz, dist.ReduceOp.MAX, row_groups[rank_c]) + # matc_nnz = matc_nnz.item() + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-nnz-{name}", barrier=True) + + # torch.cuda.nvtx.range_push("nvtx-padding") + # start_time(start_timer) + # matc_recv_indices = [] + # matc_recv_values = [] + # for i in range(replication): + # matc_recv_indices.append(torch.cuda.LongTensor(2, matc_nnz).fill_(0)) + # matc_recv_values.append(torch.cuda.DoubleTensor(matc_nnz).fill_(0.0)) + + # matc_send_indices = torch.cat((torch.cuda.LongTensor(2, matc_nnz - matc._nnz()).fill_(0), matc._indices()), 1) + # matc_send_values = torch.cat((torch.cuda.DoubleTensor(matc_nnz - matc._nnz()).fill_(0.0), matc._values())) + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-padding-{name}", barrier=True) + # torch.cuda.nvtx.range_pop() # padding + + # torch.cuda.nvtx.range_push("nvtx-allgather") + # start_time(start_timer) + # dist.all_gather(matc_recv_indices, matc_send_indices, row_groups[rank_c]) + # dist.all_gather(matc_recv_values, matc_send_values, row_groups[rank_c]) + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-allgather-{name}", barrier=True) + # torch.cuda.nvtx.range_pop() # all-gather + + # torch.cuda.nvtx.range_push("nvtx-preproc-reduce") + # start_time(start_timer) + # matc_recv = [] + # for i in range(replication): + # # matc_recv.append(torch.sparse_coo_tensor(matc_recv_indices[i], matc_recv_values[i], matc.size())) + + # # Unclear why this is necessary but it seems to hang otherwise + # nnz_mask = matc_recv_values[i] > 0 + # matc_nnz_indices = matc_recv_indices[i][:, nnz_mask] + # matc_nnz_values = matc_recv_values[i][nnz_mask] + + # matc_recv.append(torch.sparse_coo_tensor(matc_nnz_indices, matc_nnz_values, matc.size())) + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}", barrier=True) + # torch.cuda.nvtx.range_pop() # preproc-reduce + + # torch.cuda.nvtx.range_push("nvtx-reduce") + # start_time(start_timer) + # # matc_recv = torch.stack(matc_recv) + # # matc = torch.sparse.sum(matc_recv, dim=0) + + # gpu = torch.device(f"cuda:{torch.cuda.current_device()}") + # for i in range(replication): + # recv_rank = rank_c * replication + i + # if recv_rank != rank: + # # matc += matc_recv[i] + # matc_indices = matc._indices().int() + # matc_recv_indices = matc_recv[i]._indices().int() + # matc_values = matc._values().double() + + # matc_outputs = coogeam_gpu(matc_indices[0,:], matc_indices[1,:], matc_values, + # matc_recv_indices[0,:], matc_recv_indices[1,:], matc_recv[i]._values(), + # matc.size(0), matc.size(1)) + # matc_chunk_rows = matc_outputs[0].long() + # matc_chunk_cols = matc_outputs[1].long() + # matc_chunk_values = matc_outputs[2].double() + # matc_chunk_counts = torch.diff(matc_chunk_rows) + # matc_chunk_rows = torch.repeat_interleave( + # torch.arange(0, matc.size(0), device=gpu), + # matc_chunk_counts) + # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + # matc = sparse_coo_tensor_gpu(matc_chunk_indices, matc_chunk_values, matc.size()) + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) + # torch.cuda.nvtx.range_pop() # reduce + + # torch.cuda.nvtx.range_push("nvtx-reduce-coalesce") + # start_time(start_timer) + # matc = matc.coalesce() + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}", barrier=True) + # torch.cuda.nvtx.range_pop() # reduce-coalesce + + # torch.cuda.nvtx.range_push("nvtx-unpad") + # start_time(start_timer) + # # nnz_mask = matc._values() != -1.0 + # nnz_mask = matc._values() > 0 + # matc_nnz_indices = matc._indices()[:, nnz_mask] + # matc_nnz_values = matc._values()[nnz_mask] + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-unpad-{name}", barrier=True) + # torch.cuda.nvtx.range_pop() # unpad + # + # return matc_nnz_indices, matc_nnz_values + +def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication, + rank, size, row_groups, col_groups, sa_masks, + timing_dict, name, timing_arg, replicate_graph): + + global timing + timing = timing_arg + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + # TODO: assume n_layers=1 for now + # start_time(start_timer) + start_timer.record() + # p_num_indices, p_num_values = dist_saspgemm15D( + # numerator, adj_matrix, replication, rank, size, + # row_groups, col_groups, "prob", sa_masks, + # timing_dict, name) + if name == "sage" and replicate_graph: + matc_chunk_row_lens = torch.cuda.LongTensor(numerator.size(0)).fill_(0) + mata_rows = numerator._indices()[0,:] + mata_cols = numerator._indices()[1,:] + adj_row_lens = adj_matrix.crow_indices()[1:] - adj_matrix.crow_indices()[:-1] + # print(f"matc_chunk_row_lens.size: {matc_chunk_row_lens.size()}", flush=True) + # print(f"adj_row_lens.size: {adj_row_lens.size()}", flush=True) + # print(f"mata_rows.min: {mata_rows.min()} max: {mata_rows.max()}", flush=True) + # print(f"mata_cols.min: {mata_cols.min()} max: {mata_cols.max()}", flush=True) + matc_chunk_row_lens[mata_rows] = adj_row_lens[mata_cols] + # matc_chunk_crows = torch.cuda.LongTensor(numerator.size(0) + 1).fill_(0) + matc_chunk_crows = torch.cuda.IntTensor(numerator.size(0) + 1).fill_(0) + matc_chunk_crows[1:] = torch.cumsum(matc_chunk_row_lens, 0) + matc_chunk_crows[0] = 0 + matc_chunk_cols = torch.cuda.IntTensor(matc_chunk_crows[-1].item()).fill_(0) + rearrange_rows_gpu(numerator._indices()[0,:].long(), + numerator._indices()[1,:].long(), + matc_chunk_crows.long(), + adj_matrix.crow_indices().int(), + adj_matrix.col_indices().int(), + matc_chunk_cols) + matc_chunk_cols = matc_chunk_cols.long() + # p_num_rows = torch.repeat_interleave( + # torch.arange(matc_chunk_crows.size(0) - 1, + # dtype=torch.int32, + # device=numerator.device), + # matc_chunk_row_lens) + # p_num_indices = torch.stack((p_num_rows, matc_chunk_cols)) + p_num_crows = matc_chunk_crows + p_num_cols = matc_chunk_cols + # p_num_values = torch.cuda.FloatTensor(p_num_indices.size(1)).fill_(1.0) + # p_num_values = torch.cuda.DoubleTensor(p_num_indices.size(1)).fill_(1.0) + p_num_values = torch.cuda.DoubleTensor(p_num_cols.size(0)).fill_(1.0) + else: + if sa_masks is not None: + sa_masks.fill_(False) + p_num_indices, p_num_values = dist_saspgemm15D( + numerator, adj_matrix, replication, rank, size, + row_groups, col_groups, "prob", sa_masks, + timing_dict, name) + p_num_values = p_num_values.double() + + # p_num_indices, p_num_values = dist_spgemm15D(numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob") + stop_timer.record() + torch.cuda.synchronize() + time_taken = start_timer.elapsed_time(stop_timer) + # timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) + timing_dict["probability-spgemm"].append(time_taken) + + start_time(start_timer) + if name == "ladies": + # p_num_values = torch.square(p_num_values).double() + p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) + p_num_values = p_num_values.double() + p_den.scatter_add_(0, p_num_indices[0, :].long(), p_num_values) + normalize_gpu(p_num_values, p_den, p_num_indices[0, :].long(), p_num_values.size(0)) + p_num_values = torch.nan_to_num(p_num_values) + p = sparse_coo_tensor_gpu(p_num_indices.long(), p_num_values, torch.Size([numerator.size(0), node_count_total])) + elif name == "sage": + # p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) + # p_num_values = p_num_values.double() + # p_den = torch.cuda.FloatTensor(numerator.size(0)).fill_(0) + # p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) + p_num_crows = p_num_crows.long() + normalize_csr_gpu(p_num_values, p_num_crows, p_num_crows.size(0) - 1) + p_num_values = torch.nan_to_num(p_num_values) + p = torch.sparse_csr_tensor(p_num_crows, p_num_cols.long(), p_num_values, torch.Size([numerator.size(0), node_count_total])) + # p_den.scatter_add_(0, p_num_indices[0, :].long(), p_num_values) + # p = torch.sparse_coo_tensor(indices=p_num_indices, + # values=p_num_values, + # size=(numerator.size(0), node_count_total)) + # p = sparse_coo_tensor_gpu(p_num_indices, p_num_values, torch.Size([numerator.size(0), node_count_total])) + # normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) + torch.cuda.synchronize() + timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) + return p + +def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, + rank, size, row_groups, col_groups, timing_dict, name): + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + sample_start_timer = torch.cuda.Event(enable_timing=True) + sample_stop_timer = torch.cuda.Event(enable_timing=True) + + select_start_timer = torch.cuda.Event(enable_timing=True) + select_stop_timer = torch.cuda.Event(enable_timing=True) + + select_iter_start_timer = torch.cuda.Event(enable_timing=True) + select_iter_stop_timer = torch.cuda.Event(enable_timing=True) + + rank_c = rank // replication + rank_col = rank % replication + + # n_darts_col = n_darts // replication + # if rank_col == replication - 1: + # n_darts_col = n_darts - (replication - 1) * n_darts_col + + # n_darts_col = n_darts + # avg_degree = int(p._nnz() / p.size(0)) + # n_darts_col = int(avg_degree * frontier_size / avg_degree) + n_darts_col = frontier_size + + p_row_lens = p.crow_indices()[1:] - p.crow_indices()[:-1] + next_frontier_rows = torch.repeat_interleave( + torch.arange(p.crow_indices().size(0) - 1, + dtype=torch.int32, + device=p.device), + p_row_lens) + next_frontier_indices = torch.stack((next_frontier_rows, p.col_indices())) + next_frontier = torch.sparse_coo_tensor(indices=next_frontier_indices, + values=torch.cuda.LongTensor(p._nnz()).fill_(0), + size=(p.size(0), node_count_total)) + # next_frontier = sparse_coo_tensor_gpu(p._indices(), torch.cuda.LongTensor(p._nnz()).fill_(0), + # torch.Size([p.size(0), node_count_total])) + + start_time(start_timer) + + frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0)).fill_(0) + ones = torch.cuda.IntTensor(next_frontier._indices()[0,:].size(0)).fill_(1) + if name == "ladies": + frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) + zero_idxs = p._indices()[0, (p._values() == 0)].unique() + else: + sum_csri_gpu(frontier_nnz_sizes, p.crow_indices(), ones, p.crow_indices().size(0) - 1) + zero_idxs = next_frontier._indices()[0, (p.values() == 0)].unique() + zero_count = (next_frontier._indices()[0,:] == 0).nonzero().size(0) + frontier_nnz_sizes[0] = zero_count + + # frontier_nnz_sizes[zero_idxs] = 0 + frontier_nnz_sizes[zero_idxs] -= 1 + + timing_dict["sample-pre-loop"].append(stop_time(start_timer, stop_timer)) + sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) + + iter_count = 0 + selection_iter_count = 0 + + # underfull_minibatches = (sampled_count < frontier_size).any() + underfull_minibatches = True + + ps_frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0) + 1).fill_(0) + ps_frontier_nnz_sizes[1:] = torch.cumsum(frontier_nnz_sizes, dtype=torch.int32, dim=0) + p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) + # p_rowsum = torch.cuda.FloatTensor(p.size(0)).fill_(0) + while underfull_minibatches: + iter_count += 1 + start_time(sample_start_timer) + + start_time(start_timer) + + if name == "ladies": + ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) + ps_p_values[0] = 0 + p_rowsum.fill_(0) + p_rowsum.scatter_add_(0, p._indices()[0, :], p._values()) + ps_p_rowsum = torch.cumsum(p_rowsum, dim=0).roll(1) + ps_p_rowsum[0] = 0 + elif name == "sage": + ps_p_values = torch.cumsum(p.values(), dim=0) + p_rowsum.fill_(0) + # p_rowsum.scatter_add_(0, p._indices()[0, :], p._values()) + sum_csrd_gpu(p_rowsum, p.crow_indices(), p.values(), p.crow_indices().size(0) - 1) + # ps_p_rowsum = torch.cumsum(p_rowsum, dim=0) + ps_p_rowsum = torch.cumsum(p_rowsum, dim=0).roll(1) + ps_p_rowsum[0] = 0 + timing_dict["sample-prob-rowsum"].append(stop_time(start_timer, stop_timer)) + + # n_darts_col = int((frontier_size - sampled_count).sum().item() / p.size(0)) * 5 + 1 + rank_batches_start = rank_col * (sampled_count.size(0) // replication) + rank_batches_stop = (rank_col + 1) * (sampled_count.size(0) // replication) + rank_batches_stop = min(sampled_count.size(0), rank_batches_stop) + dart_count_row = frontier_size - sampled_count + dart_count_row[:rank_batches_start] = 0 + dart_count_row[rank_batches_stop:] = 0 + n_darts_col = dart_count_row.sum().item() + + if n_darts_col > 0: + start_time(start_timer) + # dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() + # dart_values = torch.cuda.DoubleTensor(n_darts_col * p.size(0)).uniform_() + # dart_values = torch.cuda.FloatTensor(n_darts_col * p.size(0)).uniform_() + # dart_values = torch.cuda.FloatTensor(n_darts_col).uniform_() + dart_values = torch.cuda.DoubleTensor(n_darts_col).uniform_() + # ps_dart_count_row = torch.cumsum(dart_count_row, dim=0, dtype=torch.int32).roll(1) + ps_dart_count_row = torch.cumsum(dart_count_row, dim=0) + timing_dict["sample-gen-darts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # compute_darts1d_gpu(dart_values, n_darts, mb_count) + # compute_darts1d_gpu(dart_values, n_darts_col, mb_count) + # compute_darts1d_gpu(dart_values, p_rowsum, ps_p_rowsum, n_darts_col, p.size(0)) + compute_darts1d_gpu(dart_values, p_rowsum, ps_p_rowsum, ps_dart_count_row, n_darts_col, p.size(0)) + + timing_dict["sample-dart-throw"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) + throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ + n_darts_col, p._nnz()) + timing_dict["sample-filter-darts"].append(stop_time(start_timer, stop_timer)) + + frontier_nnz_start = ps_frontier_nnz_sizes[rank_batches_start] + frontier_nnz_stop = ps_frontier_nnz_sizes[rank_batches_stop] + dart_hits_count[:frontier_nnz_start] = 0 + dart_hits_count[frontier_nnz_stop:] = 0 + else: + dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) + + start_time(start_timer) + if replication > 1: + dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) + next_frontier_values = torch.logical_or( + dart_hits_count, + next_frontier._values().int()).int() + # next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), + # values=next_frontier_values, + # size=(p.size(0), node_count_total)) + next_frontier_tmp = sparse_coo_tensor_gpu(next_frontier._indices(), next_frontier_values, + torch.Size([p.size(0), node_count_total])) + timing_dict["sample-add-to-frontier"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + sampled_count_old = sampled_count.clone() + sampled_count = sampled_count.fill_(0) + sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) + if name == "ladies": + next_frontier_nnzmask = next_frontier_values.nonzero().squeeze() + next_frontier_nnzvals = next_frontier_values[next_frontier_nnzmask] + next_frontier_nnzidxs = next_frontier_tmp._indices()[0, next_frontier_nnzmask] + # sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0, :], next_frontier_values) + sampled_count.scatter_add_(0, next_frontier_nnzidxs, next_frontier_nnzvals) + else: + sum_csri_gpu(sampled_count, p.crow_indices(), next_frontier_values, p.crow_indices().size(0) - 1) + timing_dict["sample-count-samples"].append(stop_time(start_timer, stop_timer)) + + del sampled_count_old + + start_time(start_timer) + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + overflowed_minibatches = (overflow > 0).any() + timing_dict["sample-select-preproc"].append(stop_time(start_timer, stop_timer)) + + start_time(select_start_timer) + + if rank_col == 0 and overflowed_minibatches: + while overflowed_minibatches: + start_time(select_iter_start_timer) + + start_time(start_timer) + selection_iter_count += 1 + ps_overflow = torch.cumsum(overflow, dim=0) + total_overflow = ps_overflow[-1].item() + # n_darts_select = total_overflow // replication + # if rank_col == replication - 1: + # n_darts_select = total_overflow - (replication - 1) * n_darts_select + n_darts_select = total_overflow + timing_dict["select-psoverflow"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + dart_hits_inv = dart_hits_count.reciprocal().double() + dart_hits_inv[dart_hits_inv == float("inf")] = 0 + timing_dict["select-reciprocal"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + dart_hits_inv_mtx = sparse_coo_tensor_gpu(next_frontier._indices(), dart_hits_inv, + torch.Size([p.size(0), node_count_total])) + timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + dart_hits_inv_sum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) + dart_hits_inv_nnzidxs = dart_hits_inv.nonzero().squeeze() + indices_nnz_idxs = next_frontier._indices()[0,dart_hits_inv_nnzidxs] + dart_hits_inv_sum.scatter_add_(0, indices_nnz_idxs, dart_hits_inv[dart_hits_inv_nnzidxs]) + timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + ps_dart_hits_inv_sum = torch.cumsum(dart_hits_inv_sum, dim=0).roll(1) + ps_dart_hits_inv_sum[0] = 0 + timing_dict["select-psinv"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # dart_select = torch.cuda.DoubleTensor(total_overflow).uniform_() + dart_select = torch.cuda.DoubleTensor(n_darts_select).uniform_() + # Compute darts for selection + compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, + p.size(0), n_darts_select) + timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # Throw selection darts + ps_dart_hits_inv = torch.cumsum(dart_hits_inv, dim=0) + # throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count, total_overflow, + throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count, n_darts_select, + next_frontier._nnz()) + timing_dict["select-throwdarts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # dist.all_reduce(dart_hits_count, op=dist.ReduceOp.MIN, group=row_groups[rank_c]) + next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).int() + # next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), + # values=next_frontier_values, + # size=(p.size(0), node_count_total)) + next_frontier_tmp = sparse_coo_tensor_gpu(next_frontier._indices(), next_frontier_values, + torch.Size([p.size(0), node_count_total])) + timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + sampled_count = torch.cuda.IntTensor(p.size(0)).fill_(0) + sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) + # sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values()) + next_frontier_nnzvals = next_frontier_tmp._values().nonzero().squeeze() + next_frontier_nnzidxs = next_frontier_tmp._indices()[0,next_frontier_nnzvals] + sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ + next_frontier_tmp._values()[next_frontier_nnzvals]) + timing_dict["select-samplecount"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + overflowed_minibatches = (overflow > 0).any() + timing_dict["select-overflow"].append(stop_time(start_timer, stop_timer)) + + timing_dict["select-iter"].append(stop_time(select_iter_start_timer, select_iter_stop_timer)) + else: + timing_dict["select-psoverflow"] = [] + timing_dict["select-reciprocal"] = [] + timing_dict["select-instmtx"] = [] + timing_dict["select-invsum"] = [] + timing_dict["select-psinv"] = [] + timing_dict["select-computedarts"] = [] + timing_dict["select-throwdarts"] = [] + timing_dict["select-add-to-frontier"] = [] + timing_dict["select-samplecount"] = [] + timing_dict["select-overflow"] = [] + timing_dict["select-iter"] = [] + + # dist.barrier(group=row_groups[rank_c]) + + timing_dict["sample-dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) + + start_time(start_timer) + if replication > 1: + dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) + dist.broadcast(dart_hits_count, src=rank_c * replication, group=row_groups[rank_c]) + next_frontier_values = torch.logical_or( + dart_hits_count, + next_frontier._values()).int() + # next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), + # values=next_frontier_values, + # size=(p.size(0), node_count_total)) + next_frontier = sparse_coo_tensor_gpu(next_frontier._indices(), next_frontier_values, + torch.Size([p.size(0), node_count_total])) + timing_dict["sample-gather-counts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + dart_hits_mask = dart_hits_count > 0 + p.values()[dart_hits_mask] = 0.0 + + filled_minibatches = sampled_count == frontier_size + # filled_minibatches_mask = torch.gather(filled_minibatches, 0, p._indices()[0,:]) + filled_minibatches_mask = torch.gather(filled_minibatches, 0, next_frontier._indices()[0,:]) + p.values()[filled_minibatches_mask] = 0 + timing_dict["sample-set-probs"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + sampled_count = torch.cuda.IntTensor(p.size(0)).fill_(0) + sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) + if name == "ladies": + next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() + next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] + sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ + next_frontier._values()[next_frontier_nnzvals].int()) + else: + sum_csri_gpu(sampled_count, p.crow_indices(), next_frontier._values(), p.crow_indices().size(0) - 1) + underfull_minibatches = (sampled_count < frontier_size).any() + timing_dict["sample-compute-bool"].append(stop_time(start_timer, stop_timer)) + + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + + timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) + + return next_frontier + +def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ + batch_size, frontier_size, mb_count, mb_count_total, node_count_total, replication, \ + rank, size, row_groups, col_groups, timing_dict, layer_id, name, semibulk_size=128): + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + semibulk_size = min(semibulk_size, mb_count // replication) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-construct-nextf") + if name == "ladies": + next_frontier_select = torch.masked_select(next_frontier._indices()[1,:], \ + next_frontier._values().bool()).view(mb_count, frontier_size) + batches_select = torch.masked_select(batches._indices()[1,:], \ + batches._values().bool()).view(mb_count, batch_size) + elif name == "sage": + next_frontier_select = next_frontier._indices()[1,:].view(mb_count * batch_size, frontier_size) + batches_select = torch.masked_select(batches._indices()[1,:], \ + batches._values().bool()).view(mb_count * batch_size, 1) + next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) + torch.cuda.nvtx.range_pop() # construct-nextf + timing_dict["construct-nextf"].append(stop_time(start_timer, stop_timer, barrier=True)) + + torch.cuda.nvtx.range_push("nvtx-select-rowcols") + # 1. Make the row/col select matrices + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-mtxs") + + if layer_id == 0: + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz * mb_count).cuda(), + batches_select.view(-1))) + else: + # TODO i > 0 (make sure size args to spspmm are right) + print("ERROR i > 0") + + row_select_mtx_values = torch.cuda.DoubleTensor(row_select_mtx_indices[0].size(0)).fill_(1.0) + + repeated_next_frontier = next_frontier_select.clone() + if name == "ladies": + scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() + elif name == "sage": + scale_mtx = torch.arange(start=0, end=node_count_total * mb_count * batch_size, + step=node_count_total).cuda() + scale_mtx = scale_mtx[:, None] + repeated_next_frontier.add_(scale_mtx) + col_select_mtx_rows = repeated_next_frontier.view(-1) + if name == "ladies": + col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).cuda().repeat( + mb_count,1).view(-1) + elif name == "sage": + col_select_mtx_cols = torch.arange(next_frontier_select.size(1) * batch_size).cuda().repeat( + mb_count,1).view(-1) + col_select_mtx_indices = torch.stack((col_select_mtx_rows, col_select_mtx_cols)) + + # col_select_mtx_values = torch.cuda.DoubleTensor(col_select_mtx_rows.size(0)).fill_(1.0) + + col_select_mtx_values = torch.cuda.DoubleTensor(col_select_mtx_rows.size(0)).fill_(0.0) + col_unique_rows, col_inverse = torch.unique(col_select_mtx_rows, sorted=True, return_inverse=True) + col_rows_perm = torch.arange(col_inverse.size(0), dtype=col_inverse.dtype, device=col_inverse.device) + col_row_mask = col_inverse.new_empty(col_unique_rows.size(0)).scatter_(0, col_inverse, col_rows_perm) + col_select_mtx_values[col_row_mask] = 1.0 + + torch.cuda.nvtx.range_pop() + timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer, barrier=True)) + + # 2. Multiply row_select matrix with adj_matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") + # row_select_mtx = torch.sparse_coo_tensor(row_select_mtx_indices, row_select_mtx_values, + # size=(nnz * mb_count, node_count_total)) + row_select_mtx = sparse_coo_tensor_gpu(row_select_mtx_indices, row_select_mtx_values, + torch.Size([nnz * mb_count, node_count_total])) + if sa_masks is not None: + sa_masks.fill_(0) + row_select_mtx = row_select_mtx.to_sparse_csr() + # adj_matrix = adj_matrix.to_sparse_csr() + sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ + row_groups, col_groups, "rowsel", sa_masks, + timing_dict) + row_select_mtx = row_select_mtx.to_sparse_coo() + # adj_matrix = adj_matrix.to_sparse_coo() + # sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + # size=(nnz * mb_count, node_count_total)) + sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + torch.Size([nnz * mb_count, node_count_total])) + + torch.cuda.nvtx.range_pop() + timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer, barrier=True)) + + # 3. Expand sampled rows + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-row-select-expand") + row_shift = torch.cuda.LongTensor(sampled_values.size(0)).fill_(0) + if name == "ladies": + # semibulk_size = mb_count // (64 * replication) # number of batches to column extract from in bulk + # semibulk_size = 4 # number of batches to column extract from in bulk + # semibulk_size = 128 # number of batches to column extract from in bulk + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, replication, batch_size, node_count_total, mb_count_total, + batch_size, semibulk_size) + elif name == "sage": + semibulk_size = mb_count * nnz # vtxs to column extract from in bulk + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, replication, batch_size, node_count_total, mb_count_total, 1, + semibulk_size) + sampled_indices[1,:] += row_shift + torch.cuda.nvtx.range_pop() + timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer, barrier=True)) + + # 4. Multiply sampled rows with col_select matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") + if name == "ladies": + # sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + # size=(nnz * mb_count, node_count_total * mb_count_total)) + sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + torch.Size([nnz * mb_count, node_count_total * semibulk_size])) + # sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + # torch.Size([nnz * mb_count, node_count_total])) + elif name == "sage": + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + size=(nnz * mb_count, node_count_total * mb_count_total * nnz)) + # sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + # torch.Size([nnz * mb_count, node_count_total * semibulk_size])) + # sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + # torch.Size([nnz * mb_count, node_count_total])) + if name == "ladies": + # col_select_mtx = torch.sparse_coo_tensor(col_select_mtx_indices, col_select_mtx_values, + # size=(node_count_total * mb_count, next_frontier.size(1))) + # col_select_mtx = sparse_coo_tensor_gpu(col_select_mtx_indices, col_select_mtx_values, + # torch.Size([node_count_total * mb_count, next_frontier.size(1)])) + col_select_mtx = sparse_coo_tensor_gpu(col_select_mtx_indices, col_select_mtx_values, + torch.Size([node_count_total * mb_count, next_frontier_select.size(1)])) + elif name == "sage": + # col_select_mtx = torch.sparse_coo_tensor(col_select_mtx_indices, col_select_mtx_values, + # size=(node_count_total * mb_count * batch_size, next_frontier.size(1) * batch_size)) + col_select_mtx = torch.sparse_coo_tensor(col_select_mtx_indices, col_select_mtx_values, + torch.Size([node_count_total * mb_count * batch_size, next_frontier_select.size(1) * batch_size])) + + # sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ + # row_groups, col_groups, "colsel", timing_dict) + + sample_mtx = sample_mtx.to_sparse_csr() + # sampled_indices, sampled_values = col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, + # replication, rank, size, row_groups, col_groups, "colsel", + # timing_dict) + if name == "ladies": + bulk_items = mb_count + bulk_size = batch_size + elif name == "sage": + bulk_items = mb_count * nnz + bulk_size = 1 + sampled_indices, sampled_values = col_select15D(sample_mtx, col_select_mtx, bulk_items, node_count_total, + bulk_size, semibulk_size, replication, rank, size, row_groups, + col_groups, "colsel", name, timing_dict) + + torch.cuda.nvtx.range_pop() + timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer, barrier=True)) + + # 5. Store sampled matrices + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-set-sample") + # adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ + # size=(nnz * mb_count, next_frontier.size(1))) + # adj_matrix_sample = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + # torch.Size([nnz * mb_count, next_frontier.size(1)])) + adj_matrix_sample = sparse_coo_tensor_gpu(sampled_indices, sampled_values, + torch.Size([nnz * mb_count, next_frontier_select.size(1)])) + + torch.cuda.nvtx.range_pop() + timing_dict["set-sample"].append(stop_time(start_timer, stop_timer, barrier=True)) + + torch.cuda.nvtx.range_pop() # nvtx-select-mtxs + + return batches_select, next_frontier_select, adj_matrix_sample + +class GraphDataset(Dataset): + """ + The custom default GNN dataset to load graphs off the disk + """ + + def __init__( + self, + input_dir, + data_name=None, + num_events=None, + stage="fit", + hparams=None, + transform=None, + pre_transform=None, + pre_filter=None, + preprocess=True, + ): + if hparams is None: + hparams = {} + super().__init__(input_dir, transform, pre_transform, pre_filter) + + self.input_dir = input_dir + self.data_name = data_name + self.hparams = hparams + self.num_events = num_events + self.stage = stage + self.preprocess = preprocess + self.transform = transform + + self.input_paths = load_datafiles_in_dir( + self.input_dir, self.data_name, self.num_events + ) + self.input_paths.sort() # We sort here for reproducibility + + def len(self): + return len(self.input_paths) + + def get(self, idx): + event_path = self.input_paths[idx] + print(f"event_path: {event_path}", flush=True) + # if "event005000555.pyg" in event_path: + # return None + event_name = event_path.split("/")[-1] + event_id = int(event_name[5:event_name.index(".")]) + + # if "CTD" in event_path: + # event = torch.load(f"/global/cfs/cdirs/m1982/alokt/data/trackml/ctd/trainset_processed/{event_name}", map_location=torch.device("cpu")) + # return event + + event = torch.load(event_path, map_location=torch.device("cpu")) + # convert DataBatch to Data instance because some transformations don't work on DataBatch + event = Data(**event.to_dict()) + if not self.preprocess: + return event + event = self.preprocess_event(event) + # do pyg transformation if a torch_geometric.transforms instance is given + if self.transform is not None: + event = self.transform(event) + + # return (event, event_path) if self.stage == "predict" else event + return event + + def preprocess_event(self, event): + """ + Process event before it is used in training and validation loops + """ + event = self.apply_hard_cuts(event) + event = self.construct_weighting(event) + event = self.handle_edge_list(event) + event = self.scale_features(event) + if self.hparams.get("edge_features") is not None: + event = self.add_edge_features( + event + ) # scaling must be done before adding features + return event + + def apply_hard_cuts(self, event): + """ + Apply hard cuts to the event. This is implemented by + 1. Finding which true edges are from tracks that pass the hard cut. + 2. Pruning the input graph to only include nodes that are connected to these edges. + """ + + if ( + self.hparams is not None + and "hard_cuts" in self.hparams.keys() + and self.hparams["hard_cuts"] + ): + assert isinstance( + self.hparams["hard_cuts"], dict + ), "Hard cuts must be a dictionary" + handle_hard_cuts(event, self.hparams["hard_cuts"]) + + return event + + def construct_weighting(self, event): + """ + Construct the weighting for the event + """ + + assert event.y.shape[0] == event.edge_index.shape[1], ( + f"Input graph has {event.edge_index.shape[1]} edges, but" + f" {event.y.shape[0]} truth labels" + ) + + if self.hparams is not None and "weighting" in self.hparams.keys(): + assert isinstance(self.hparams["weighting"], list) & isinstance( + self.hparams["weighting"][0], dict + ), "Weighting must be a list of dictionaries" + event.weights = handle_weighting(event, self.hparams["weighting"]) + else: + event.weights = torch.ones_like(event.y, dtype=torch.float32) + + return event + + def handle_edge_list(self, event): + if ( + "input_cut" in self.hparams.keys() + and self.hparams["input_cut"] + and "scores" in event.keys() + ): + # Apply a score cut to the event + self.apply_score_cut(event, self.hparams["input_cut"]) + + # if "undirected" in self.hparams.keys() and self.hparams["undirected"]: + # # Flip event.edge_index and concat together + # self.to_undirected(event) + return event + + def to_undirected(self, event): + """ + Add the reverse of the edge_index to the event. This then requires all edge features to be duplicated. + Additionally, the truth map must be duplicated. + """ + num_edges = event.edge_index.shape[1] + # Flip event.edge_index and concat together + event.edge_index = torch.cat( + [event.edge_index, event.edge_index.flip(0)], dim=1 + ) + # event.edge_index, unique_edge_indices = torch.unique(event.edge_index, dim=1, return_inverse=True) + num_track_edges = event.track_edges.shape[1] + event.track_edges = torch.cat( + [event.track_edges, event.track_edges.flip(0)], dim=1 + ) + + # Concat all edge-like features together + for key in get_pyg_data_keys(event): + if key == "truth_map": + continue + if not isinstance(event[key], torch.Tensor) or not event[key].shape: + continue + if event[key].shape[0] == num_edges: + event[key] = torch.cat([event[key], event[key]], dim=0) + elif event[key].shape[0] == num_track_edges: + event[key] = torch.cat([event[key], event[key]], dim=0) + + # handle truth_map separately + truth_map = event.truth_map.clone() + truth_map[truth_map >= 0] = truth_map[truth_map >= 0] + num_edges + event.truth_map = torch.cat([event.truth_map, truth_map], dim=0) + + return event + + def add_edge_features(self, event): + if "edge_features" in self.hparams.keys(): + assert isinstance( + self.hparams["edge_features"], list + ), "Edge features must be a list of strings" + handle_edge_features(event, self.hparams["edge_features"]) + return event + + def scale_features(self, event): + """ + Handle feature scaling for the event + """ + + if ( + self.hparams is not None + and "node_scales" in self.hparams.keys() + and "node_features" in self.hparams.keys() + ): + assert isinstance( + self.hparams["node_scales"], list + ), "Feature scaling must be a list of ints or floats" + for i, feature in enumerate(self.hparams["node_features"]): + assert feature in get_pyg_data_keys( + event + ), f"Feature {feature} not found in event" + event[feature] = event[feature] / self.hparams["node_scales"][i] + + return event + + def unscale_features(self, event): + """ + Unscale features when doing prediction + """ + + if ( + self.hparams is not None + and "node_scales" in self.hparams.keys() + and "node_features" in self.hparams.keys() + ): + assert isinstance( + self.hparams["node_scales"], list + ), "Feature scaling must be a list of ints or floats" + for i, feature in enumerate(self.hparams["node_features"]): + assert feature in get_pyg_data_keys( + event + ), f"Feature {feature} not found in event" + event[feature] = event[feature] * self.hparams["node_scales"][i] + return event + + def apply_score_cut(self, event, score_cut): + """ + Apply a score cut to the event. This is used for the evaluation stage. + """ + passing_edges_mask = event.scores >= score_cut + num_edges = event.edge_index.shape[1] + for key in get_pyg_data_keys(event): + if ( + isinstance(event[key], torch.Tensor) + and event[key].shape + and ( + event[key].shape[0] == num_edges + or event[key].shape[-1] == num_edges + ) + ): + event[key] = event[key][..., passing_edges_mask] + + remap_from_mask(event, passing_edges_mask) + return event + + def get_y_node(self, event): + y_node = torch.zeros(event.z.size(0)) + y_node[event.track_edges.view(-1)] = 1 + event.y_node = y_node + return event + +def load_datafiles_in_dir(input_dir, data_name=None, data_num=None): + if data_name is not None: + input_dir = os.path.join(input_dir, data_name) + + data_files = [str(path) for path in Path(input_dir).rglob("*.pyg")][:data_num] + print(f"data_files: {data_files}", flush=True) + assert len(data_files) > 0, f"No data files found in {input_dir}" + if data_num is not None: + assert len(data_files) == data_num, ( + f"Number of data files found ({len(data_files)}) is less than the number" + f" requested ({data_num})" + ) + + return data_files + +def handle_hard_cuts(event, hard_cuts_config): + true_track_mask = torch.ones_like(event.truth_map, dtype=torch.bool) + + for condition_key, condition_val in hard_cuts_config.items(): + assert condition_key in get_pyg_data_keys( + event + ), f"Condition key {condition_key} not found in event keys {get_pyg_data_keys(event)}" + condition_lambda = get_condition_lambda(condition_key, condition_val) + value_mask = condition_lambda(event) + true_track_mask = true_track_mask * value_mask + + graph_mask = torch.isin( + event.edge_index, event.track_edges[:, true_track_mask] + ).all(0) + remap_from_mask(event, graph_mask) + + num_edges = event.edge_index.shape[1] + for edge_key in get_pyg_data_keys(event): + if ( + isinstance(event[edge_key], torch.Tensor) + and num_edges in event[edge_key].shape + ): + event[edge_key] = event[edge_key][..., graph_mask] + + num_track_edges = event.track_edges.shape[1] + for track_feature in get_pyg_data_keys(event): + if ( + isinstance(event[track_feature], torch.Tensor) + and num_track_edges in event[track_feature].shape + ): + event[track_feature] = event[track_feature][..., true_track_mask] + +def reset_angle(angles): + angles[angles > torch.pi] = angles[angles > torch.pi] - 2 * torch.pi + angles[angles < -torch.pi] = angles[angles < -torch.pi] + 2 * torch.pi + return angles + +def handle_edge_features(event, edge_features): + src, dst = event.edge_index + + for edge_feature in edge_features: + if "dr" in edge_features and not ("dr" in get_pyg_data_keys(event)): + event.dr = event.r[dst] - event.r[src] + if "dphi" in edge_features and not ("dphi" in get_pyg_data_keys(event)): + event.dphi = ( + reset_angle((event.phi[dst] - event.phi[src]) * torch.pi) / torch.pi + ) + if "dz" in edge_features and not ("dz" in get_pyg_data_keys(event)): + event.dz = event.z[dst] - event.z[src] + if "deta" in edge_features and not ("deta" in get_pyg_data_keys(event)): + event.deta = event.eta[dst] - event.eta[src] + if "phislope" in edge_features and not ("phislope" in get_pyg_data_keys(event)): + dr = event.r[dst] - event.r[src] + dphi = reset_angle((event.phi[dst] - event.phi[src]) * torch.pi) / torch.pi + phislope = dphi / dr + event.phislope = phislope + if "phislope" in edge_features: + event.phislope = torch.nan_to_num( + event.phislope, nan=0.0, posinf=100, neginf=-100 + ) + event.phislope = torch.clamp(event.phislope, -100, 100) + if "rphislope" in edge_features and not ( + "rphislope" in get_pyg_data_keys(event) + ): + r_ = (event.r[dst] + event.r[src]) / 2.0 + dr = event.r[dst] - event.r[src] + dphi = reset_angle((event.phi[dst] - event.phi[src]) * torch.pi) / torch.pi + phislope = dphi / dr + phislope = torch.nan_to_num(phislope, nan=0.0, posinf=100, neginf=-100) + phislope = torch.clamp(phislope, -100, 100) + rphislope = torch.multiply(r_, phislope) + event.rphislope = rphislope # features / norm / pre_proc once + if "rphislope" in edge_features: + event.rphislope = torch.nan_to_num(event.rphislope, nan=0.0) + +def get_pyg_data_keys(event: Data): + """ + Get the keys of the pyG data object. + """ + if torch_geometric.__version__ < "2.4.0": + return event.keys + else: + return event.keys() + +def get_condition_lambda(condition_key, condition_val): + condition_dict = { + "is": lambda event: event[condition_key] == condition_val, + "is_not": lambda event: event[condition_key] != condition_val, + "in": lambda event: torch.isin( + event[condition_key], + torch.tensor(condition_val[1], device=event[condition_key].device), + ), + "not_in": lambda event: ~torch.isin( + event[condition_key], + torch.tensor(condition_val[1], device=event[condition_key].device), + ), + "within": lambda event: (condition_val[0] <= event[condition_key].float()) + & (event[condition_key].float() <= condition_val[1]), + "not_within": lambda event: not ( + (condition_val[0] <= event[condition_key].float()) + & (event[condition_key].float() <= condition_val[1]) + ), + } + + if isinstance(condition_val, bool): + return lambda event: event[condition_key] == condition_val + elif isinstance(condition_val, list) and not isinstance(condition_val[0], str): + return lambda event: (condition_val[0] <= event[condition_key].float()) & ( + event[condition_key].float() <= condition_val[1] + ) + elif isinstance(condition_val, list): + return condition_dict[condition_val[0]] + else: + raise ValueError(f"Condition {condition_val} not recognised") + +def remap_from_mask(event, edge_mask): + """ + Takes a mask applied to the edge_index tensor, and remaps the truth_map tensor indices to match. + """ + + truth_map_to_edges = torch.ones(edge_mask.shape[0], dtype=torch.long) * -1 + truth_map_to_edges = truth_map_to_edges.to(event.truth_map.device) + truth_map_to_edges[event.truth_map[event.truth_map >= 0]] = torch.arange( + event.truth_map.shape[0], device=event.truth_map.device + )[event.truth_map >= 0] + truth_map_to_edges = truth_map_to_edges[edge_mask] + + new_map = torch.ones(event.truth_map.shape[0], dtype=torch.long) * -1 + new_map = new_map.to(event.truth_map.device) + new_map[truth_map_to_edges[truth_map_to_edges >= 0]] = torch.arange( + truth_map_to_edges.shape[0], device=truth_map_to_edges.device + )[truth_map_to_edges >= 0] + event.truth_map = new_map.to(event.truth_map.device) + +def handle_weighting(event, weighting_config): + """ + Take the specification of the weighting and convert this into float values. The default is: + - True edges have weight 1.0 + - Negative edges have weight 1.0 + + The weighting_config can be used to change this behaviour. For example, we might up-weight target particles - that is edges that pass: + - y == 1 + - primary == True + - pt > 1 GeV + - etc. As desired. + + We can also down-weight (i.e. mask) edges that are true, but not of interest. For example, we might mask: + - y == 1 + - primary == False + - pt < 1 GeV + - etc. As desired. + """ + + # Set the default values, which will be overwritten if specified in the config + weights = torch.zeros_like(event.y, dtype=torch.float) + weights[event.y == 0] = 1.0 + + for weight_spec in weighting_config: + weight_val = weight_spec["weight"] + weights[get_weight_mask(event, weight_spec["conditions"])] = weight_val + + return weights + +def get_weight_mask(event, weight_conditions): + graph_mask = torch.ones_like(event.y) + + for condition_key, condition_val in weight_conditions.items(): + assert condition_key in get_pyg_data_keys( + event + ), f"Condition key {condition_key} not found in event keys {get_pyg_data_keys(event)}" + condition_lambda = get_condition_lambda(condition_key, condition_val) + value_mask = condition_lambda(event) + graph_mask = graph_mask * map_tensor_handler( + value_mask, + output_type="edge-like", + num_nodes=event.num_nodes, + edge_index=event.edge_index, + truth_map=event.truth_map, + ) + + return graph_mask + +def map_tensor_handler( + input_tensor: torch.Tensor, + output_type: str, + input_type: str = None, + truth_map: torch.Tensor = None, + edge_index: torch.Tensor = None, + track_edges: torch.Tensor = None, + num_nodes: int = None, + num_edges: int = None, + num_track_edges: int = None, + aggr: str = None, +): + """ + A general function to handle arbitrary maps of one tensor type to another + Types are "node-like", "edge-like" and "track-like". + - node-like: The input tensor is of the same size as the + number of nodes in the graph + - edge-like: The input tensor is of the same size as the + number of edges in the graph, that is, the *constructed* graph + - track-like: The input tensor is of the same size as the + number of true track edges in the event, that is, the *truth* graph + + To visualize: + (n) + ^ + / \ + edge_to_node / \ track_to_node + / \ + / \ + / \ + / \ + / \ +node_to_edge / \ node_to_track + / \ + v edge_to_track v + (e) <-------------> (t) + track_to_edge + + Args: + input_tensor (torch.Tensor): The input tensor to be mapped + output_type (str): The type of the output tensor. + One of "node-like", "edge-like" or "track-like" + input_type (str, optional): The type of the input tensor. + One of "node-like", "edge-like" or "track-like". Defaults to None, + and will try to infer the type from the input tensor, if num_nodes + and/or num_edges are provided. + truth_map (torch.Tensor, optional): The truth map tensor. + Defaults to None. Used for mappings to/from track-like tensors. + num_nodes (int, optional): The number of nodes in the graph. + Defaults to None. Used for inferring the input type. + num_edges (int, optional): The number of edges in the graph. + Defaults to None. Used for inferring the input type. + num_track_edges (int, optional): The number of track edges in the graph + Defaults to None. Used for inferring the input type. + """ + + if num_track_edges is None and truth_map is not None: + num_track_edges = truth_map.shape[0] + if num_track_edges is None and track_edges is not None: + num_track_edges = track_edges.shape[1] + if num_edges is None and edge_index is not None: + num_edges = edge_index.shape[1] + if input_type is None: + input_type, input_tensor = infer_input_type( + input_tensor, num_nodes, num_edges, num_track_edges + ) + if input_type == output_type: + return input_tensor + + input_args = { + "truth_map": truth_map, + "edge_index": edge_index, + "track_edges": track_edges, + "num_nodes": num_nodes, + "num_edges": num_edges, + "num_track_edges": num_track_edges, + "aggr": aggr, + } + + mapping_functions = { + ("node-like", "edge-like"): map_nodes_to_edges, + ("edge-like", "node-like"): map_edges_to_nodes, + ("node-like", "track-like"): map_nodes_to_tracks, + ("track-like", "node-like"): map_tracks_to_nodes, + ("edge-like", "track-like"): map_edges_to_tracks, + ("track-like", "edge-like"): map_tracks_to_edges, + } + if (input_type, output_type) not in mapping_functions: + raise ValueError(f"Mapping from {input_type} to {output_type} not supported") + + return mapping_functions[(input_type, output_type)](input_tensor, **input_args) + +# Returns string and tensor +def infer_input_type( + input_tensor: torch.Tensor, + num_nodes: int = None, + num_edges: int = None, + num_track_edges: int = None, +) -> (str, torch.Tensor): + """ + Infers the type of the input tensor based on its shape and the provided number of nodes, edges, and track edges. + + Args: + input_tensor (torch.Tensor): The tensor whose type needs to be inferred. + num_nodes (int, optional): Number of nodes in the graph. + num_edges (int, optional): Number of edges in the graph. + num_track_edges (int, optional): Number of track edges in the graph. + + Returns: + str: The inferred type of the input tensor. One of ["node-like", "edge-like", "track-like"]. + """ + + NODE_LIKE = "node-like" + EDGE_LIKE = "edge-like" + TRACK_LIKE = "track-like" + + if num_nodes is not None and input_tensor.shape[0] == num_nodes: + return NODE_LIKE, input_tensor + elif num_edges is not None and num_edges in input_tensor.shape: + return EDGE_LIKE, input_tensor + elif num_track_edges is not None and num_track_edges in input_tensor.shape: + return TRACK_LIKE, input_tensor + elif num_track_edges is not None and num_track_edges // 2 in input_tensor.shape: + return TRACK_LIKE, torch.cat([input_tensor, input_tensor], dim=0) + else: + raise ValueError("Unable to infer the type of the input tensor.") + +def map_nodes_to_edges( + nodelike_input: torch.Tensor, edge_index: torch.Tensor, aggr: str = None, **kwargs +): + """ + Map a node-like tensor to an edge-like tensor. If the aggregation is None, this is simply done by sending node values to the edges, thus returning a tensor of shape (2, num_edges). + If the aggregation is not None, the node values are aggregated to the edges, and the resulting tensor is of shape (num_edges,). + """ + + if aggr is None: + return nodelike_input[edge_index] + + edgelike_tensor = nodelike_input[edge_index] + torch_aggr = getattr(torch, aggr) + return torch_aggr(edgelike_tensor, dim=0) + +def map_edges_to_nodes( + edgelike_input: torch.Tensor, + edge_index: torch.Tensor, + aggr: str = None, + num_nodes: int = None, + **kwargs, +): + """ + Map an edge-like tensor to a node-like tensor. If the aggregation is None, this is simply done by sending edge values to the nodes, thus returning a tensor of shape (num_nodes,). + If the aggregation is not None, the edge values are aggregated to the nodes at the destination node (edge_index[1]), and the resulting tensor is of shape (num_nodes,). + """ + + if num_nodes is None: + num_nodes = int(edge_index.max().item() + 1) + + if aggr is None: + nodelike_output = torch.zeros( + num_nodes, dtype=edgelike_input.dtype, device=edgelike_input.device + ) + nodelike_output[edge_index] = edgelike_input + return nodelike_output + + return scatter( + edgelike_input, edge_index[1], dim=0, dim_size=num_nodes, reduce=aggr + ) + +def map_nodes_to_tracks( + nodelike_input: torch.Tensor, track_edges: torch.Tensor, aggr: str = None, **kwargs +): + """ + Map a node-like tensor to a track-like tensor. If the aggregation is None, this is simply done by sending node values to the tracks, thus returning a tensor of shape (2, num_track_edges). + If the aggregation is not None, the node values are aggregated to the tracks, and the resulting tensor is of shape (num_track_edges,). + """ + + if aggr is None: + return nodelike_input[track_edges] + + tracklike_tensor = nodelike_input[track_edges] + torch_aggr = getattr(torch, aggr) + return torch_aggr(tracklike_tensor, dim=0) + +def map_tracks_to_nodes( + tracklike_input: torch.Tensor, + track_edges: torch.Tensor, + aggr: str = None, + num_nodes: int = None, + **kwargs, +): + """ + Map a track-like tensor to a node-like tensor. If the aggregation is None, this is simply done by sending track values to the nodes, thus returning a tensor of shape (num_nodes,). + If the aggregation is not None, the track values are aggregated to the nodes at the destination node (track_edges[1]), and the resulting tensor is of shape (num_nodes,). + """ + + if num_nodes is None: + num_nodes = int(track_edges.max().item() + 1) + + if aggr is None: + nodelike_output = torch.zeros( + num_nodes, dtype=tracklike_input.dtype, device=tracklike_input.device + ) + nodelike_output[track_edges] = tracklike_input + return nodelike_output + + return scatter( + tracklike_input.repeat(2), + torch.cat([track_edges[0], track_edges[1]]), + dim=0, + dim_size=num_nodes, + reduce=aggr, + ) + +def map_tracks_to_edges( + tracklike_input: torch.Tensor, + truth_map: torch.Tensor, + num_edges: int = None, + **kwargs, +): + """ + Map an track-like tensor to a edge-like tensor. This is done by sending the track value through the truth map, where the truth map is >= 0. Note that where truth_map == -1, + the true edge has not been constructed in the edge_index. In that case, the value is set to NaN. + """ + + if num_edges is None: + num_edges = int(truth_map.max().item() + 1) + edgelike_output = torch.zeros( + num_edges, dtype=tracklike_input.dtype, device=tracklike_input.device + ) + edgelike_output[truth_map[truth_map >= 0]] = tracklike_input[truth_map >= 0] + edgelike_output[truth_map[truth_map == -1]] = float("nan") + return edgelike_output + + +def map_edges_to_tracks( + edgelike_input: torch.Tensor, truth_map: torch.Tensor, **kwargs +): + """ + TODO: Implement this. I don't think it is a meaningful operation, but it is needed for completeness. + """ + raise NotImplementedError( + "This is not a meaningful operation, but it is needed for completeness" + ) + +def rearrange_by_distance(event, edge_index): + assert "r" in get_pyg_data_keys(event) and "z" in get_pyg_data_keys( + event + ), "event must contain r and z" + distance = event.r**2 + event.z**2 + + # flip edges that are pointing inward + edge_mask = distance[edge_index[0]] > distance[edge_index[1]] + edge_index[:, edge_mask] = edge_index[:, edge_mask].flip(0) + + return edge_index + +def graph_roc_curve(dataset_name, test_batches, title, filename): + """ + Plot the ROC curve for the graph construction efficiency. + """ + print( + "Plotting the ROC curve and score distribution, events from" + f" {dataset_name}" + ) + all_y_truth, all_scores, masked_scores, masked_y_truth = [], [], [], [] + masks = [] + # dataset_name = config["dataset"] + # dataset = getattr(lightning_module, dataset_name) + + # for event in tqdm(dataset): + for event in tqdm(test_batches): + # event = event.to(lightning_module.device) + event = event.cuda() + # Need to apply score cut and remap the truth_map + if "weights" in get_pyg_data_keys(event): + target_y = event.weights.bool() & event.y.bool() + mask = event.weights > 0 + else: + target_y = event.y.bool() + mask = torch.ones_like(target_y).bool().to(target_y.device) + + all_y_truth.append(target_y) + all_scores.append(event.scores) + masked_scores.append(event.scores[mask]) + masked_y_truth.append(target_y[mask]) + masks.append(mask) + + all_scores = torch.cat(all_scores).cpu().numpy() + all_y_truth = torch.cat(all_y_truth).cpu().numpy() + masked_scores = torch.cat(masked_scores).cpu().numpy() + masked_y_truth = torch.cat(masked_y_truth).cpu().numpy() + masks = torch.cat(masks).cpu().numpy() + + fig, ax = plt.subplots(figsize=(8, 6)) + # Get the ROC curve + fpr, tpr, _ = roc_curve(all_y_truth, all_scores) + full_auc_score = auc(fpr, tpr) + + # Plot the ROC curve + ax.plot(fpr, tpr, color="black", label="ROC curve") + + # Get the ROC curve + fpr, tpr, _ = roc_curve(masked_y_truth, masked_scores) + masked_auc_score = auc(fpr, tpr) + + # Plot the ROC curve + ax.plot(fpr, tpr, color="green", label="masked ROC curve") + + ax.plot([0, 1], [0, 1], color="black", linestyle="--", label="Random classifier") + ax.set_xlabel("False Positive Rate", ha="right", x=0.95, fontsize=14) + ax.set_ylabel("True Positive Rate", ha="right", y=0.95, fontsize=14) + ax.set_xscale("log") + ax.set_yscale("log") + ax.legend(loc="lower right", fontsize=14) + ax.text( + 0.95, + 0.20, + f"Full AUC: {full_auc_score:.3f}, Masked AUC: {masked_auc_score: .3f}", + ha="right", + va="bottom", + transform=ax.transAxes, + fontsize=14, + ) + + # Save the plot + atlasify( + "Internal", + # f"{plot_config['title']} \n" + f"{title} \n" + r"$\sqrt{s}=14$TeV, $t \bar{t}$, $\langle \mu \rangle = 200$, primaries $t" + r" \bar{t}$ and soft interactions) " + "\n" + r"$p_T > 1$GeV, $|\eta| < 4$" + + "\n" + # + f"Evaluated on {dataset.len()} events in {dataset_name}", + + f"Evaluated on {len(test_batches)} events in {dataset_name}", + ) + # filename_template = plot_config.get("filename") + filename_template = filename + filename = ( + f"{filename_template}_roc_curve.png" + if filename_template is not None + else "roc_curve.png" + ) + # filename = os.path.join(config["stage_dir"], filename) + filename = os.path.join(".", filename) + fig.savefig(filename) + print("Finish plotting. Find the ROC curve at" f" {filename}") + plt.close() + fig, ax = plt.subplots(figsize=(8, 6)) + all_y_truth = all_y_truth.astype(np.int16) + all_y_truth[~masks] = 2 + labels = np.array(["Fake"] * len(all_y_truth)) + labels[all_y_truth == 1] = "Target True" + labels[all_y_truth == 2] = "Non-target True" + # weight = 1 / dataset.len() + weight = 1 / len(test_batches) + ax = plot_score_histogram(all_scores, labels, ax=ax, inverse_dataset_length=weight) + ax.set_xlabel("Edge score", ha="right", x=0.95, fontsize=14) + ax.set_ylabel("Count/event", ha="right", y=0.95, fontsize=14) + atlasify( + "Internal", + "Score Distribution \n" + r"$\sqrt{s}=14$TeV, $t \bar{t}$, $\langle \mu \rangle = 200$, primaries $t" + r" \bar{t}$ and soft interactions) " + "\n" + r"$p_T > 1$GeV, $|\eta| < 4$" + + "\n" + # + f"Evaluated on {dataset.len()} events in {dataset_name}", + + f"Evaluated on {len(test_batches)} events in {dataset_name}", + ) + filename = ( + f"{filename_template}_score_distribution.png" + if filename_template is not None + else "score_distribution.png" + ) + # filename = os.path.join(config["stage_dir"], filename) + filename = os.path.join(".", filename) + fig.savefig(filename) + print("Finish plotting. Find the score distribution at" f" {filename}") + + return full_auc_score, masked_auc_score + +def graph_intersection( + input_pred_graph, + input_truth_graph, + return_y_pred=True, + return_y_truth=False, + return_pred_to_truth=False, + return_truth_to_pred=False, + unique_pred=True, + unique_truth=True, +): + """ + An updated version of the graph intersection function, which is around 25x faster than the + Scipy implementation (on GPU). Takes a prediction graph and a truth graph, assumed to have unique entries. + If unique_pred or unique_truth is False, the function will first find the unique entries in the input graphs, and return the updated edge lists. + """ + + if not unique_pred: + input_pred_graph = torch.unique(input_pred_graph, dim=1) + if not unique_truth: + input_truth_graph = torch.unique(input_truth_graph, dim=1) + + unique_edges, inverse = torch.unique( + torch.cat([input_pred_graph, input_truth_graph], dim=1), + dim=1, + sorted=False, + return_inverse=True, + return_counts=False, + ) + + inverse_pred_map = torch.ones_like(unique_edges[1]) * -1 + inverse_pred_map[inverse[: input_pred_graph.shape[1]]] = torch.arange( + input_pred_graph.shape[1], device=input_pred_graph.device + ) + + inverse_truth_map = torch.ones_like(unique_edges[1]) * -1 + inverse_truth_map[inverse[input_pred_graph.shape[1] :]] = torch.arange( + input_truth_graph.shape[1], device=input_truth_graph.device + ) + + pred_to_truth = inverse_truth_map[inverse][: input_pred_graph.shape[1]] + truth_to_pred = inverse_pred_map[inverse][input_pred_graph.shape[1] :] + + return_tensors = [] + + if not unique_pred: + return_tensors.append(input_pred_graph) + if not unique_truth: + return_tensors.append(input_truth_graph) + if return_y_pred: + y_pred = pred_to_truth >= 0 + return_tensors.append(y_pred) + if return_y_truth: + y_truth = truth_to_pred >= 0 + return_tensors.append(y_truth) + if return_pred_to_truth: + return_tensors.append(pred_to_truth) + if return_truth_to_pred: + return_tensors.append(truth_to_pred) + + return return_tensors if len(return_tensors) > 1 else return_tensors[0] + +def apply_score_cut(event, score_cut): + """ + Apply a score cut to the event. This is used for the evaluation stage. + """ + passing_edges_mask = event.scores >= score_cut + + # flip edge direction if points inward + event.edge_index = rearrange_by_distance(event, event.edge_index) + event.track_edges = rearrange_by_distance(event, event.track_edges) + + event.graph_truth_map = graph_intersection( + event.edge_index, + event.track_edges, + return_y_pred=False, + return_y_truth=False, + return_truth_to_pred=True, + ) + event.truth_map = graph_intersection( + event.edge_index[:, passing_edges_mask], + event.track_edges, + return_y_pred=False, + return_truth_to_pred=True, + ) + event.pred = passing_edges_mask + +def apply_target_conditions(event, target_tracks): + """ + Apply the target conditions to the event. This is used for the evaluation stage. + Target_tracks is a list of dictionaries, each of which contains the conditions to be applied to the event. + """ + # passing_tracks = torch.ones(event.truth_map.shape[0], dtype=torch.bool).to( + # self.device + # ) + passing_tracks = torch.ones(event.truth_map.shape[0], dtype=torch.bool).cuda() + + for condition_key, condition_val in target_tracks.items(): + condition_lambda = get_condition_lambda(condition_key, condition_val) + # passing_tracks = passing_tracks * condition_lambda(event).to(self.device) + passing_tracks = passing_tracks * condition_lambda(event).cuda() + + event.target_mask = passing_tracks + +def gnn_efficiency_rz(test_batches, plot_config: dict, config: dict): + """_summary_ + + Args: + plot_config (dict): any plotting config + config (dict): config + + Plot GNN edgewise efficiency against rz + """ + + print("Plotting edgewise efficiency as a function of rz") + print( + f"Using score cut: {config.get('score_cut')}, events from {config['dataset']}" + ) + if "target_tracks" in config: + print(f"Track selection criteria: \n{yaml.dump(config.get('target_tracks'))}") + else: + print("No track selection criteria found, accepting all tracks.") + + target = {"z": torch.empty(0), "r": torch.empty(0)} + all_target = target.copy() + true_positive = target.copy() + input_graph_size, graph_size, n_graphs = (0, 0, 0) + + # dataset_name = config["dataset"] + # dataset = getattr(lightning_module, dataset_name) + + # for event in tqdm(dataset): + for event in tqdm(test_batches): + # event = event.to(lightning_module.device) + event = event.cuda() + print(f"event.r: {event.r}", flush=True) + print(f"event.z: {event.z}", flush=True) + + # Need to apply score cut and remap the truth_map + if "score_cut" in config: + # lightning_module.apply_score_cut(event, config["score_cut"]) + apply_score_cut(event, config["score_cut"]) + if "target_tracks" in config: + apply_target_conditions(event, config["target_tracks"]) + else: + event.target_mask = torch.ones(event.truth_map.shape[0], dtype=torch.bool) + + # scale r and z + event.r /= 1000 + event.z /= 1000 + + # flip edges that point inward if not undirected, since if undirected is True, lightning_module.apply_score_cut takes care of this + event.edge_index = rearrange_by_distance(event, event.edge_index) + event.track_edges = rearrange_by_distance(event, event.track_edges) + event = event.cpu() + + # indices of all target edges present in the input graph + print(f"event.target_mask: {event.target_mask}", flush=True) + print(f"event.graph_truth_map: {event.graph_truth_map}", flush=True) + target_edges = event.track_edges[ + :, event.target_mask & (event.graph_truth_map > -1) + ] + + # indices of all target edges (may or may not be present in the input graph) + all_target_edges = event.track_edges[:, event.target_mask] + + # get target z r + for key, item in target.items(): + target[key] = torch.cat([item, event[key][target_edges[0]]], dim=0) + for key, item in all_target.items(): + all_target[key] = torch.cat([item, event[key][all_target_edges[0]]], dim=0) + + # indices of all true positive target edges + target_true_positive_edges = event.track_edges[ + :, event.target_mask & (event.truth_map > -1) + ] + for key in ["r", "z"]: + true_positive[key] = torch.cat( + [true_positive[key], event[key][target_true_positive_edges[0]]], dim=0 + ) + + input_graph_size += event.edge_index.size(1) + graph_size += event.pred.sum().numpy() + n_graphs += 1 + event.r *= 1000 # rescale because events test_batches are maintained across function calls + event.z *= 1000 + + print(f"plot_efficiency_rz1", flush=True) + fig, ax = plot_efficiency_rz( + target["z"], target["r"], true_positive["z"], true_positive["r"], plot_config + ) + # Save the plot + atlasify( + "Internal", + r"$\sqrt{s}=14$TeV, $t \bar{t}$, $\langle \mu \rangle = 200$, primaries $t" + r" \bar{t}$ and soft interactions) " + "\n" + r"$p_T > 1$ GeV, $ | \eta | < 4$ " + "\n" + "Graph Construction Efficiency:" + f" {(target['z'].shape[0] / all_target['z'].shape[0]):.4f}, Input graph size:" + f" {input_graph_size / n_graphs: .2e} \n" + r"Edge score cut: " + + str(config["score_cut"]) + + f", Mean graph size: {graph_size / n_graphs :.2e} \n" + "Signal Efficiency:" + f" {true_positive['z'].shape[0] / target['z'].shape[0] :.4f} \n" + "Cumulative signal efficiency:" + f" {true_positive['z'].shape[0] / all_target['z'].shape[0]: .4f}" + + "\n" + # + f"Evaluated on {dataset.len()} events in {dataset_name}", + + f"Evaluated on {len(test_batches)} events", + ) + plt.tight_layout() + save_dir = os.path.join( + f"{plot_config.get('filename', 'edgewise_efficiency_rz')}.png", + ) + fig.savefig(save_dir) + print(f"Finish plotting. Find the plot at {save_dir}") + plt.close() + + print(f"plot_efficiency_rz2", flush=True) + fig, ax = plot_efficiency_rz( + all_target["z"], + all_target["r"], + true_positive["z"], + true_positive["r"], + plot_config, + ) + # Save the plot + atlasify( + "Internal", + r"$\sqrt{s}=14$TeV, $t \bar{t}$, $\langle \mu \rangle = 200$, primaries $t" + r" \bar{t}$ and soft interactions) " + "\n" + r"$p_T > 1$ GeV, $ | \eta | < 4$ " + "\n" + "Graph Construction Efficiency:" + f" {(target['z'].shape[0] / all_target['z'].shape[0]):.4f}, Input graph size:" + f" {input_graph_size / n_graphs: .2e} \n" + r"Edge score cut: " + + str(config["score_cut"]) + + f", Mean graph size: {graph_size / n_graphs :.2e} \n" + "Signal Efficiency:" + f" {true_positive['z'].shape[0] / target['z'].shape[0] :.4f} \n" + "Cumulative signal efficiency:" + f" {true_positive['z'].shape[0] / all_target['z'].shape[0]: .4f}" + + "\n" + # + f"Evaluated on {dataset.len()} events in {dataset_name}", + + f"Evaluated on {len(test_batches)} events", + ) + plt.tight_layout() + save_dir = os.path.join( + f"cumulative_{plot_config.get('filename', 'edgewise_efficiency_rz')}.png", + ) + fig.savefig(save_dir) + print(f"Finish plotting. Find the plot at {save_dir}") + plt.close() + + signal_efficiency = true_positive['z'].shape[0] / target['z'].shape[0] + return signal_efficiency + +def gnn_purity_rz(test_batches, plot_config: dict, config: dict): + """_summary_ + + Args: + plot_config (dict): any plotting config + config (dict): config + + Plot GNN edgewise efficiency against rz + """ + + print("Plotting edgewise purity as a function of rz") + print( + f"Using score cut: {config.get('score_cut')}, events from {config['dataset']}" + ) + if "target_tracks" in config: + print(f"Track selection criteria: \n{yaml.dump(config.get('target_tracks'))}") + else: + print("No track selection criteria found, accepting all tracks.") + + true_positive = { + # key: torch.empty(0).to(lightning_module.device) for key in ["z", "r"] + key: torch.empty(0).cuda() for key in ["z", "r"] + } + target_true_positive = true_positive.copy() + + pred = true_positive.copy() + masked_pred = true_positive.copy() + + # dataset_name = config["dataset"] + # dataset = getattr(lightning_module, dataset_name) + + # for event in tqdm(dataset): + for event in tqdm(test_batches): + # event = event.to(lightning_module.device) + event = event.cuda() + # Need to apply score cut and remap the truth_map + if "score_cut" in config: + apply_score_cut(event, config["score_cut"]) + if "target_tracks" in config: + apply_target_conditions(event, config["target_tracks"]) + else: + event.target_mask = torch.ones(event.truth_map.shape[0], dtype=torch.bool) + + # scale r and z + event.r /= 1000 + event.z /= 1000 + print(f"event: {event}", flush=True) + print(f"event.r: {event.r}", flush=True) + print(f"event.z: {event.z}", flush=True) + + # flip edges that point inward if not undirected, since if undirected is True, lightning_module.apply_score_cut takes care of this + event.edge_index = rearrange_by_distance(event, event.edge_index) + event.track_edges = rearrange_by_distance(event, event.track_edges) + # event = event.cpu() + + # target true positive edge indices, used as numerator of target purity and purity + target_true_positive_edges = event.track_edges[ + :, event.target_mask & (event.truth_map > -1) + ] + + # true positive edge indices, used as numerator of total purity + true_positive_edges = event.track_edges[:, (event.truth_map > -1)] + + # all positive edges, used as denominator of total and target purity + positive_edges = event.edge_index[:, event.pred] + + # masked positive edge indices, including true positive target edges and all false positive edges + fake_positive_edges = event.edge_index[:, event.pred & (event.y == 0)] + masked_positive_edges = torch.cat( + [target_true_positive_edges, fake_positive_edges], dim=1 + ) + + for key in ["r", "z"]: + target_true_positive[key] = torch.cat( + [ + target_true_positive[key].float(), + event[key][target_true_positive_edges[0]].float(), + ], + dim=0, + ) + true_positive[key] = torch.cat( + [ + true_positive[key].float(), + event[key][true_positive_edges[0]].float(), + ], + dim=0, + ) + pred[key] = torch.cat( + [pred[key].float(), event[key][positive_edges[0]].float()], dim=0 + ) + masked_pred[key] = torch.cat( + [ + masked_pred[key].float(), + event[key][masked_positive_edges[0]].float(), + ], + dim=0, + ) + print(f"true_positive[z]: {true_positive['z']}", flush=True) + print(f"true_positive[r]: {true_positive['r']}", flush=True) + event.r *= 1000 # rescale because events test_batches are maintained across function calls + event.z *= 1000 + + purity_definition_label = { + "target_purity": "Target Purity", + "masked_purity": "Masked Purity", + "total_purity": "Total Purity", + } + for numerator, denominator, suffix in zip( + [true_positive, target_true_positive, target_true_positive], + [pred, pred, masked_pred], + ["total_purity", "target_purity", "masked_purity"], + ): + print(f"plot_efficiency_rz3", flush=True) + print(f"numerator[z]: {numerator['z']}", flush=True) + print(f"numerator[r]: {numerator['r']}", flush=True) + fig, ax = plot_efficiency_rz( + denominator["z"].cpu(), + denominator["r"].cpu(), + numerator["z"].cpu(), + numerator["r"].cpu(), + plot_config, + ) + # Save the plot + atlasify( + "Internal", + r"$\sqrt{s}=14$TeV, $t \bar{t}$, $\langle \mu \rangle = 200$, primaries $t" + r" \bar{t}$ and soft interactions) " + "\n" + r"$p_T > 1$ GeV, $ | \eta | < 4$" + "\n" + r"Edge score cut: " + + str(config["score_cut"]) + + "\n" + + purity_definition_label[suffix] + + ": " + + f"{numerator['z'].size(0) / denominator['z'].size(0) : .5f}" + + "\n" + # + f"Evaluated on {dataset.len()} events in {dataset_name}", + + f"Evaluated on {len(test_batches)} events", + ) + plt.tight_layout() + save_dir = os.path.join( + f"{plot_config.get('filename', 'edgewise')}_{suffix}_rz.png", + ) + fig.savefig(save_dir) + print(f"Finish plotting. Find the plot at {save_dir}") + plt.close() + +def plot_efficiency_rz( + target_z: torch.Tensor, + target_r: torch.Tensor, + true_positive_z: torch.Tensor, + true_positive_r: torch.Tensor, + plot_config: dict, +): + z_range, r_range = plot_config.get("z_range", [-3, 3]), plot_config.get( + "r_range", [0, 1.0] + ) + z_bins, r_bins = plot_config.get("z_bins", 6 * 64), plot_config.get("r_bins", 64) + z_bins = np.linspace(z_range[0], z_range[1], z_bins, endpoint=True) + r_bins = np.linspace(r_range[0], r_range[1], r_bins, endpoint=True) + + fig, ax = plt.subplots(1, 1, figsize=plot_config.get("fig_size", (12, 6))) + true_hist, _, _ = np.histogram2d( + target_z.numpy(), + target_r.numpy(), + bins=[z_bins, r_bins], + ) + true_positive_hist, z_edges, r_edges = np.histogram2d( + true_positive_z.numpy(), true_positive_r.numpy(), bins=[z_bins, r_bins] + ) + + print(f"target_z: {target_z} z_bins: {z_bins}", flush=True) + print(f"target_r: {target_r} r_bins: {r_bins}", flush=True) + print(f"true_positive_hist: {true_positive_hist}", flush=True) + print(f"true_hist: {true_hist}", flush=True) + eff = true_positive_hist / true_hist + print(f"eff.T: {eff.T}", flush=True) + print(f"eff.T.nonnans: {np.count_nonzero(~np.isnan(eff.T))} eff.T.numelems: {eff.T.size}", flush=True) + + c = ax.pcolormesh( + z_bins, + r_bins, + eff.T, + cmap="jet_r", + vmin=plot_config.get("vmin", 0.9), + vmax=plot_config.get("vmax", 1), + ) + fig.colorbar(c, ax=ax) + ax.set_xlabel("z [m]") + ax.set_ylabel("r [m]") + + return fig, ax + +def plot_score_histogram(scores, y, bins=100, ax=None, inverse_dataset_length=1): + """ + Plot a histogram of scores, labelled by truth + """ + # weight each entry by the inverse of dataset length, such that the y axis relects the count per event per bin + weights = np.array([inverse_dataset_length] * len(scores)) + ax = sns.histplot( + x=scores, + hue=y, + bins=100, + stat="count", + weights=weights, + log_scale=(False, True), + ax=ax, + element="step", + palette="colorblind", + fill=False, + ) + sns.move_legend( + ax, + "lower center", + bbox_to_anchor=(0.5, 1), + ncol=np.unique(y).shape[0], + title=None, + frameon=False, + ) + + return ax + +# class GraphDataset(Dataset): +# """ +# The custom default GNN dataset to load graphs off the disk +# """ +# +# def __init__(self, input_dir, data_name = None, num_events = None, stage="fit", hparams=None, transform=None, pre_transform=None, pre_filter=None): +# super().__init__(input_dir, transform, pre_transform, pre_filter) +# +# self.input_dir = input_dir +# self.data_name = data_name +# self.hparams = hparams +# self.num_events = num_events +# self.stage = stage +# +# self.input_paths = load_datafiles_in_dir(self.input_dir, self.data_name, self.num_events) +# self.input_paths.sort() # We sort here for reproducibility +# +# def len(self): +# return len(self.input_paths) +# +# def get(self, idx): +# +# event_path = self.input_paths[idx] +# event = torch.load(event_path, map_location=torch.device("cpu")) +# #del event.track_edges +# #event.track_edges=torch.t(event.track_edges) +# #print("Event: ", event) +# self.preprocess_event(event) +# #print(self.hparams["node_features"]) ['r', 'phi', 'z'] +# #event["x"] = torch.stack([event[feature] for feature in self.hparams["node_features"]], dim=-1).float() +# #event = Data(x=event['x'], edge_index=event['edge_index'],track_edges=event["track_edges"], truth_map=event["truth_map"], y=event["y"]) +# #print("Event:", event_path, event) +# # return (event, event_path) if self.stage == "predict" else event +# return event +# +# def preprocess_event(self, event): +# """ +# Process event before it is used in training and validation loops +# """ +# +# print(f"preprocessing event", flush=True) +# self.apply_hard_cuts(event) +# self.construct_weighting(event) +# self.handle_edge_list(event) +# self.add_edge_features(event) +# self.scale_features(event) +# +# def apply_hard_cuts(self, event): +# """ +# Apply hard cuts to the event. This is implemented by +# 1. Finding which true edges are from tracks that pass the hard cut. +# 2. Pruning the input graph to only include nodes that are connected to these edges. +# """ +# +# if self.hparams is not None and "hard_cuts" in self.hparams.keys() and self.hparams["hard_cuts"]: +# assert isinstance(self.hparams["hard_cuts"], dict), "Hard cuts must be a dictionary" +# handle_hard_cuts(event, self.hparams["hard_cuts"]) +# +# def construct_weighting(self, event): +# """ +# Construct the weighting for the event +# """ +# +# assert event.y.shape[0] == event.edge_index.shape[1], f"Input graph has {event.edge_index.shape[1]} edges, but {event.y.shape[0]} truth labels" +# +# if self.hparams is not None and "weighting" in self.hparams.keys(): +# assert isinstance(self.hparams["weighting"], list) & isinstance(self.hparams["weighting"][0], dict), "Weighting must be a list of dictionaries" +# event.weights = handle_weighting(event, self.hparams["weighting"]) +# else: +# event.weights = torch.ones_like(event.y, dtype=torch.float32) +# +# def handle_edge_list(self, event): +# +# if "input_cut" in self.hparams.keys() and self.hparams["input_cut"] and "scores" in event.keys: +# # Apply a score cut to the event +# self.apply_score_cut(event, self.hparams["input_cut"]) +# +# # if "undirected" in self.hparams.keys() and self.hparams["undirected"]: +# # # Flip event.edge_index and concat together +# # self.to_undirected(event) +# +# +# def to_undirected(self, event): +# """ +# Add the reverse of the edge_index to the event. This then requires all edge features to be duplicated. +# Additionally, the truth map must be duplicated. +# """ +# +# num_edges = event.edge_index.shape[1] +# # Flip event.edge_index and concat together +# event.edge_index = torch.cat([event.edge_index, event.edge_index.flip(0)], dim=1) +# # event.edge_index, unique_edge_indices = torch.unique(event.edge_index, dim=1, return_inverse=True) +# +# # Concat all edge-like features together +# for key in event.keys: +# if isinstance(event[key], torch.Tensor) and ((event[key].shape[0] == num_edges)): +# event[key] = torch.cat([event[key], event[key]], dim=0) +# # event[key] = torch.zeros_like(event.edge_index[0], dtype=event[key].dtype).scatter(0, unique_edge_indices, event[key]) +# +# +# def add_edge_features(self, event): +# if "edge_features" in self.hparams.keys(): +# assert isinstance(self.hparams["edge_features"], list), "Edge features must be a list of strings" +# handle_edge_features(event, self.hparams["edge_features"]) +# +# def scale_features(self, event): +# """ +# Handle feature scaling for the event +# """ +# +# print(f"hparams: {self.hparams}", flush=True) +# if self.hparams is not None and "node_scales" in self.hparams.keys() and "node_features" in self.hparams.keys(): +# assert isinstance(self.hparams["node_scales"], list), "Feature scaling must be a list of ints or floats" +# for i, feature in enumerate(self.hparams["node_features"]): +# assert feature in event.keys, f"Feature {feature} not found in event" +# event[feature] = event[feature] / self.hparams["node_scales"][i] +# +# def unscale_features(self, event): +# """ +# Unscale features when doing prediction +# """ +# +# if self.hparams is not None and "node_scales" in self.hparams.keys() and "node_features" in self.hparams.keys(): +# assert isinstance(self.hparams["node_scales"], list), "Feature scaling must be a list of ints or floats" +# for i, feature in enumerate(self.hparams["node_features"]): +# assert feature in event.keys, f"Feature {feature} not found in event" +# event[feature] = event[feature] * self.hparams["node_scales"][i] +# +# def apply_score_cut(self, event, score_cut): +# """ +# Apply a score cut to the event. This is used for the evaluation stage. +# """ +# passing_edges_mask = event.scores >= score_cut +# num_edges = event.edge_index.shape[1] +# for key in event.keys: +# if isinstance(event[key], torch.Tensor) and event[key].shape and (event[key].shape[0] == num_edges or event[key].shape[-1] == num_edges): +# event[key] = event[key][..., passing_edges_mask] +# +# remap_from_mask(event, passing_edges_mask) +# +# def load_datafiles_in_dir(input_dir, data_name = None, data_num = None): +# +# if data_name is not None: +# input_dir = os.path.join(input_dir, data_name) +# +# data_files = [str(path) for path in Path(input_dir).rglob("*.pyg")][:data_num] +# assert len(data_files) > 0, f"No data files found in {input_dir}" +# if data_num is not None: +# assert len(data_files) == data_num, f"Number of data files found ({len(data_files)}) is less than the number requested ({data_num})" +# +# return data_files +# +# def handle_hard_cuts(event, hard_cuts_config): +# +# true_track_mask = torch.ones_like(event.truth_map, dtype=torch.bool) +# +# print(f"config: {hard_cuts_config}", flush=True) +# for condition_key, condition_val in hard_cuts_config.items(): +# assert condition_key in event.keys, f"Condition key {condition_key} not found in event keys {event.keys}" +# condition_lambda = get_condition_lambda(condition_key, condition_val) +# value_mask = condition_lambda(event) +# true_track_mask = true_track_mask * value_mask +# +# graph_mask = torch.isin(event.edge_index, event.track_edges[:, true_track_mask]).all(0) +# remap_from_mask(event, graph_mask) +# +# num_edges = event.edge_index.shape[1] +# for edge_key in event.keys: +# if isinstance(event[edge_key], torch.Tensor) and num_edges in event[edge_key].shape: +# event[edge_key] = event[edge_key][..., graph_mask] +# +# num_track_edges = event.track_edges.shape[1] +# for track_feature in event.keys: +# if isinstance(event[track_feature], torch.Tensor) and num_track_edges in event[track_feature].shape: +# event[track_feature] = event[track_feature][..., true_track_mask] +# +# def get_condition_lambda(condition_key, condition_val): +# +# condition_dict = { +# "is": lambda event: event[condition_key] == condition_val, +# "is_not": lambda event: event[condition_key] != condition_val, +# "in": lambda event: torch.isin(event[condition_key], torch.tensor(condition_val[1], device=event[condition_key].device)), +# "not_in": lambda event: ~torch.isin(event[condition_key], torch.tensor(condition_val[1], device=event[condition_key].device)), +# "within": lambda event: (condition_val[0] <= event[condition_key].float()) & (event[condition_key].float() <= condition_val[1]), +# "not_within": lambda event: not ((condition_val[0] <= event[condition_key].float()) & (event[condition_key].float() <= condition_val[1])), +# } +# +# if isinstance(condition_val, bool): +# return lambda event: event[condition_key] == condition_val +# elif isinstance(condition_val, list) and not isinstance(condition_val[0], str): +# return lambda event: (condition_val[0] <= event[condition_key].float()) & (event[condition_key].float() <= condition_val[1]) +# elif isinstance(condition_val, list): +# return condition_dict[condition_val[0]] +# else: +# raise ValueError(f"Condition {condition_val} not recognised") +# +# def remap_from_mask(event, edge_mask): +# """ +# Takes a mask applied to the edge_index tensor, and remaps the truth_map tensor indices to match. +# """ +# +# truth_map_to_edges = torch.ones(edge_mask.shape[0], dtype=torch.long) * -1 +# truth_map_to_edges[event.truth_map[event.truth_map >= 0]] = torch.arange(event.truth_map.shape[0])[event.truth_map >= 0] +# truth_map_to_edges = truth_map_to_edges[edge_mask] +# +# new_map = torch.ones(event.truth_map.shape[0], dtype=torch.long) * -1 +# new_map[truth_map_to_edges[truth_map_to_edges >= 0]] = torch.arange(truth_map_to_edges.shape[0])[truth_map_to_edges >= 0] +# event.truth_map = new_map.to(event.truth_map.device) +# +# def handle_weighting(event, weighting_config): +# """ +# Take the specification of the weighting and convert this into float values. The default is: +# - True edges have weight 1.0 +# - Negative edges have weight 1.0 +# +# The weighting_config can be used to change this behaviour. For example, we might up-weight target particles - that is edges that pass: +# - y == 1 +# - primary == True +# - pt > 1 GeV +# - etc. As desired. +# +# We can also down-weight (i.e. mask) edges that are true, but not of interest. For example, we might mask: +# - y == 1 +# - primary == False +# - pt < 1 GeV +# - etc. As desired. +# """ +# +# # Set the default values, which will be overwritten if specified in the config +# weights = torch.zeros_like(event.y, dtype=torch.float) +# weights[event.y == 0] = 1.0 +# +# for weight_spec in weighting_config: +# weight_val = weight_spec["weight"] +# weights[get_weight_mask(event, weight_spec["conditions"])] = weight_val +# +# return weights +# +# def get_weight_mask(event, weight_conditions): +# +# graph_mask = torch.ones_like(event.y) +# +# for condition_key, condition_val in weight_conditions.items(): +# assert condition_key in event.keys, f"Condition key {condition_key} not found in event keys {event.keys}" +# condition_lambda = get_condition_lambda(condition_key, condition_val) +# value_mask = condition_lambda(event) +# graph_mask = graph_mask * map_tensor_handler(value_mask, output_type="edge-like", num_nodes = event.num_nodes, edge_index = event.edge_index, truth_map = event.truth_map) +# +# return graph_mask +# +# def map_tensor_handler(input_tensor: torch.Tensor, +# output_type: str, +# input_type: str = None, +# truth_map: torch.Tensor = None, +# edge_index: torch.Tensor = None, +# track_edges: torch.Tensor = None, +# num_nodes: int = None, +# num_edges: int = None, +# num_track_edges: int = None, +# aggr: str = None): +# """ +# A general function to handle arbitrary maps of one tensor type to another. Types are "node-like", "edge-like" and "track-like". +# - Node-like: The input tensor is of the same size as the number of nodes in the graph +# - Edge-like: The input tensor is of the same size as the number of edges in the graph, that is, the *constructed* graph +# - Track-like: The input tensor is of the same size as the number of true track edges in the event, that is, the *truth* graph +# +# To visualize: +# (n) +# ^ +# / \ +# edge_to_node / \ track_to_node +# / \ +# / \ +# / \ +# / \ +# / \ +# node_to_edge / \ node_to_track +# / \ +# | | +# v edge_to_track v +# (e) <-------------> (t) +# track_to_edge +# +# Args: +# input_tensor (torch.Tensor): The input tensor to be mapped +# output_type (str): The type of the output tensor. One of "node-like", "edge-like" or "track-like" +# input_type (str, optional): The type of the input tensor. One of "node-like", "edge-like" or "track-like". Defaults to None, and will try to infer the type from the input tensor, if num_nodes and/or num_edges are provided. +# truth_map (torch.Tensor, optional): The truth map tensor. Defaults to None. Used for mappings to/from track-like tensors. +# num_nodes (int, optional): The number of nodes in the graph. Defaults to None. Used for inferring the input type. +# num_edges (int, optional): The number of edges in the graph. Defaults to None. Used for inferring the input type. +# num_track_edges (int, optional): The number of track edges in the graph. Defaults to None. Used for inferring the input type. +# """ +# +# # Refactor the above switch case into a dictionary +# mapping_dict = { +# ("node-like", "edge-like"): lambda input_tensor, truth_map, edge_index, track_index, num_nodes, num_edges, num_track_edges, aggr: map_nodes_to_edges(input_tensor, edge_index, aggr), +# ("edge-like", "node-like"): lambda input_tensor, truth_map, edge_index, track_index, num_nodes, num_edges, num_track_edges, aggr: map_edges_to_nodes(input_tensor, edge_index, aggr, num_nodes), +# ("node-like", "track-like"): lambda input_tensor, truth_map, edge_index, track_index, num_nodes, num_edges, num_track_edges, aggr: map_nodes_to_tracks(input_tensor, track_edges, aggr), +# ("track-like", "node-like"): lambda input_tensor, truth_map, edge_index, track_index, num_nodes, num_edges, num_track_edges, aggr: map_tracks_to_nodes(input_tensor, track_edges, aggr, num_nodes), +# ("edge-like", "track-like"): lambda input_tensor, truth_map, edge_index, track_index, num_nodes, num_edges, num_track_edges, aggr: map_edges_to_tracks(input_tensor, truth_map), +# ("track-like", "edge-like"): lambda input_tensor, truth_map, edge_index, track_index, num_nodes, num_edges, num_track_edges, aggr: map_tracks_to_edges(input_tensor, truth_map, num_edges), +# } +# +# if num_track_edges is None and truth_map is not None: +# num_track_edges = truth_map.shape[0] +# if num_track_edges is None and track_edges is not None: +# num_track_edges = track_edges.shape[1] +# if num_edges is None and edge_index is not None: +# num_edges = edge_index.shape[1] +# if input_type is None: +# input_type, input_tensor = infer_input_type(input_tensor, num_nodes, num_edges, num_track_edges) +# +# if input_type == output_type: +# return input_tensor +# elif (input_type, output_type) in mapping_dict: +# return mapping_dict[(input_type, output_type)](input_tensor, truth_map, edge_index, track_edges, num_nodes, num_edges, num_track_edges, aggr) +# else: +# raise ValueError(f"Mapping from {input_type} to {output_type} not supported") +# +# def infer_input_type(input_tensor: torch.Tensor, num_nodes: int = None, num_edges: int = None, num_track_edges: int = None): +# """ +# Tries to infer the input type from the input tensor and the number of nodes, edges and track-edges in the graph. +# If the input tensor cannot be matched to any of the provided types, it is assumed to be node-like. +# """ +# +# if num_nodes is not None and input_tensor.shape[0] == num_nodes: +# return "node-like", input_tensor +# elif num_edges is not None and num_edges in input_tensor.shape: +# return "edge-like", input_tensor +# elif num_track_edges is not None and num_track_edges in input_tensor.shape: +# return "track-like", input_tensor +# elif num_track_edges is not None and num_track_edges//2 in input_tensor.shape: +# return "track-like", torch.cat([input_tensor, input_tensor], dim=0) +# else: +# return "node-like", input_tensor +# +# def map_tracks_to_edges(tracklike_input: torch.Tensor, truth_map: torch.Tensor, num_edges: int = None): +# """ +# Map an track-like tensor to a edge-like tensor. This is done by sending the track value through the truth map, where the truth map is >= 0. Note that where truth_map == -1, +# the true edge has not been constructed in the edge_index. In that case, the value is set to NaN. +# """ +# +# if num_edges is None: +# num_edges = int(truth_map.max().item() + 1) +# edgelike_output = torch.zeros(num_edges, dtype=tracklike_input.dtype, device=tracklike_input.device) +# edgelike_output[truth_map[truth_map >= 0]] = tracklike_input[truth_map >= 0] +# edgelike_output[truth_map[truth_map == -1]] = float("nan") +# return edgelike_output diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/examples/cagnet_exps.sh b/examples/cagnet_exps.sh new file mode 100755 index 0000000..2c70366 --- /dev/null +++ b/examples/cagnet_exps.sh @@ -0,0 +1,29 @@ +declare -a procarr=(1 2 4) +# declare -a bsarr=(128 256 512) +declare -a bsarr=(512) +declare -a snarr=(10 20 25) +# declare -a mbarr=(64 128 256) +declare -a mbarr=(64 128) +declare -a dataarr=("reddit" "ogbn-products") + +for bs in "${bsarr[@]}" +do + for sn in "${snarr[@]}" + do + for mb in "${mbarr[@]}" + do + for data in "${dataarr[@]}" + do + for proc in "${procarr[@]}" + do + echo $bs $sn $mb $data $proc + srun -l -n $proc --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset $data --batch-size $bs --samp-num $sn --n-bulkmb $mb --replication 1 --sample-method sage --timing --baseline &> cagnet_outputs/$data/cagnet_$data\_bs$bs\_sn$sn\_mb$mb\_p$proc\_timing.out + srun -l -n $proc --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset $data --batch-size $bs --samp-num $sn --n-bulkmb $mb --replication 1 --sample-method sage --baseline &> cagnet_outputs/$data/cagnet_$data\_bs$bs\_sn$sn\_mb$mb\_p$proc\.out + done + + srun -l -n 4 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset $data --batch-size $bs --samp-num $sn --n-bulkmb $mb --replication 2 --sample-method sage --timing --baseline &> cagnet_outputs/$data/cagnet_$data\_bs$bs\_sn$sn\_mb$mb\_p4c2_timing.out + srun -l -n 4 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset $data --batch-size $bs --samp-num $sn --n-bulkmb $mb --replication 2 --sample-method sage --baseline &> cagnet_outputs/$data/cagnet_$data\_bs$bs\_sn$sn\_mb$mb\_p4c2.out + done + done + done +done diff --git a/examples/cagnet_sageconv_torchviz b/examples/cagnet_sageconv_torchviz new file mode 100644 index 0000000..343ab02 --- /dev/null +++ b/examples/cagnet_sageconv_torchviz @@ -0,0 +1,38 @@ +digraph { + graph [size="12,12"] + node [align=left fontname=monospace fontsize=10 height=0.2 ranksep=0.1 shape=box style=filled] + 140047140675392 [label=" + ()" fillcolor=darkolivegreen1] + 140047214459104 [label=NllLossBackward0] + 140047214459344 -> 140047214459104 + 140047214459344 [label=SliceBackward0] + 140047214459296 -> 140047214459344 + 140047214459296 [label=LogSoftmaxBackward0] + 140047214459392 -> 140047214459296 + 140047214459392 [label=MmBackward0] + 140047214459440 -> 140047214459392 + 140047214459440 [label=DivBackward0] + 140047214459872 -> 140047214459440 + 140047214459872 [label=ScatterAddBackward0] + 140047214459680 -> 140047214459872 + 140047214459680 [label=IndexSelectBackward0] + 140047214459584 -> 140047214459680 + 140047214459584 [label=ReluBackward0] + 140047214489552 -> 140047214459584 + 140047214489552 [label=MmBackward0] + 140047214488160 -> 140047214489552 + 140047214488160 [label=TBackward0] + 140047214488016 -> 140047214488160 + 140047214455040 [label="layers.0.lin_l.weight + (256, 128)" fillcolor=lightblue] + 140047214455040 -> 140047214488016 + 140047214488016 [label=AccumulateGrad] + 140047214459920 -> 140047214459392 + 140047214459920 [label=TBackward0] + 140047214458960 -> 140047214459920 + 140047214456560 [label="layers.1.lin_l.weight + (24, 256)" fillcolor=lightblue] + 140047214456560 -> 140047214458960 + 140047214458960 [label=AccumulateGrad] + 140047214459104 -> 140047140675392 +} diff --git a/examples/cagnet_torchviz b/examples/cagnet_torchviz new file mode 100644 index 0000000..3bedb7d --- /dev/null +++ b/examples/cagnet_torchviz @@ -0,0 +1,38 @@ +digraph { + graph [size="12,12"] + node [align=left fontname=monospace fontsize=10 height=0.2 ranksep=0.1 shape=box style=filled] + 140236170673584 [label=" + ()" fillcolor=darkolivegreen1] + 140243124960368 [label=NllLossBackward0] + 140243124958256 -> 140243124960368 + 140243124958256 [label=LogSoftmaxBackward0] + 140243125484416 -> 140243124958256 + 140243125484416 [label=GCNFuncNONEBackward] + 140243124960656 -> 140243125484416 + 140243124960656 [label=ReluBackward0] + 140243125484656 -> 140243124960656 + 140243125484656 [label=GCNFuncNONEBackward] + 140243124958448 -> 140243125484656 + 140243124957072 [label="layers.0.weight + (128, 256)" fillcolor=lightblue] + 140243124957072 -> 140243124958448 + 140243124958448 [label=AccumulateGrad] + 140236170673104 -> 140243125484656 + 140236170673104 [label=" + (51200, 128)" fillcolor=orange] + 140243124957072 -> 140243125484656 + 140243124957072 [label="layers.0.weight + (128, 256)" fillcolor=orange] + 140243124960272 -> 140243125484416 + 140243125080208 [label="layers.1.weight + (256, 24)" fillcolor=lightblue] + 140243125080208 -> 140243124960272 + 140243124960272 [label=AccumulateGrad] + 140236170673184 -> 140243125484416 + 140236170673184 [label=" + (5120, 256)" fillcolor=orange] + 140243125080208 -> 140243125484416 + 140243125080208 [label="layers.1.weight + (256, 24)" fillcolor=orange] + 140243124960368 -> 140236170673584 +} diff --git a/examples/cagnet_torchviz.svg b/examples/cagnet_torchviz.svg new file mode 100644 index 0000000..0b57611 --- /dev/null +++ b/examples/cagnet_torchviz.svg @@ -0,0 +1,164 @@ + + + + + + + + + +139891500084416 + + () + + + +139898402179152 + +NllLossBackward0 + + + +139898402179152->139891500084416 + + + + + +139898402179728 + +LogSoftmaxBackward0 + + + +139898402179728->139898402179152 + + + + + +139898402708352 + +GCNFuncNONEBackward + + + +139898402708352->139898402179728 + + + + + +139898402178720 + +ReluBackward0 + + + +139898402178720->139898402708352 + + + + + +139898402708592 + +GCNFuncNONEBackward + + + +139898402708592->139898402178720 + + + + + +139898402179440 + +AccumulateGrad + + + +139898402179440->139898402708592 + + + + + +139898402176512 + +layers.0.weight + (128, 256) + + + +139898402176512->139898402708592 + + + + + +139898402176512->139898402179440 + + + + + +139891500083936 + + (51200, 128) + + + +139891500083936->139898402708592 + + + + + +139898402179248 + +AccumulateGrad + + + +139898402179248->139898402708352 + + + + + +139898402176592 + +layers.1.weight + (256, 24) + + + +139898402176592->139898402708352 + + + + + +139898402176592->139898402179248 + + + + + +139891500084016 + + (5120, 256) + + + +139891500084016->139898402708352 + + + + + diff --git a/examples/distributed_sampling_pyg.py b/examples/distributed_sampling_pyg.py new file mode 100644 index 0000000..d2a7f8e --- /dev/null +++ b/examples/distributed_sampling_pyg.py @@ -0,0 +1,136 @@ +import os + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F +from torch import Tensor +from torch.nn.parallel import DistributedDataParallel + +from torch_geometric.datasets import Reddit +from torch_geometric.loader import NeighborLoader +from torch_geometric.nn import SAGEConv + + +class SAGE(torch.nn.Module): + def __init__(self, in_channels: int, hidden_channels: int, + out_channels: int, num_layers: int = 2): + super().__init__() + + self.convs = torch.nn.ModuleList() + self.convs.append(SAGEConv(in_channels, hidden_channels)) + for _ in range(num_layers - 2): + self.convs.append(SAGEConv(hidden_channels, hidden_channels)) + self.convs.append(SAGEConv(hidden_channels, out_channels)) + + def forward(self, x: Tensor, edge_index: Tensor) -> Tensor: + for i, conv in enumerate(self.convs): + x = conv(x, edge_index) + if i < len(self.convs) - 1: + x = x.relu() + x = F.dropout(x, p=0.5, training=self.training) + return x + + +@torch.no_grad() +def test(loader, model, rank): + model.eval() + + total_correct = total_examples = 0 + for i, batch in enumerate(loader): + out = model(batch.x, batch.edge_index.to(rank)) + pred = out[:batch.batch_size].argmax(dim=-1) + y = batch.y[:batch.batch_size].to(rank) + total_correct += int((pred == y).sum()) + total_examples += batch.batch_size + return torch.tensor(total_correct / total_examples, device=rank) + + +def run(rank, world_size, dataset): + os.environ['MASTER_ADDR'] = 'localhost' + os.environ['MASTER_PORT'] = '12355' + dist.init_process_group('nccl', rank=rank, world_size=world_size) + + data = dataset[0] + data = data.to(rank, 'x', 'y') # Move to device for faster feature fetch. + + # Split indices into `world_size` many chunks: + train_idx = data.train_mask.nonzero(as_tuple=False).view(-1) + train_idx = train_idx.split(train_idx.size(0) // world_size)[rank] + val_idx = data.val_mask.nonzero(as_tuple=False).view(-1) + val_idx = val_idx.split(val_idx.size(0) // world_size)[rank] + test_idx = data.val_mask.nonzero(as_tuple=False).view(-1) + test_idx = test_idx.split(test_idx.size(0) // world_size)[rank] + + kwargs = dict( + data=data, + batch_size=1024, + num_neighbors=[25, 10], + drop_last=True, + num_workers=4, + persistent_workers=True, + ) + train_loader = NeighborLoader( + input_nodes=train_idx, + shuffle=True, + **kwargs, + ) + val_loader = NeighborLoader( + input_nodes=val_idx, + shuffle=False, + **kwargs, + ) + test_loader = NeighborLoader( + input_nodes=test_idx, + shuffle=False, + **kwargs, + ) + + torch.manual_seed(12345) + model = SAGE(dataset.num_features, 256, dataset.num_classes).to(rank) + model = DistributedDataParallel(model, device_ids=[rank]) + optimizer = torch.optim.Adam(model.parameters(), lr=0.001) + + for epoch in range(1, 21): + model.train() + for batch in train_loader: + print(f"batch: {batch}", flush=True) + optimizer.zero_grad() + out = model(batch.x, batch.edge_index.to(rank))[:batch.batch_size] + loss = F.cross_entropy(out, batch.y[:batch.batch_size]) + loss.backward() + optimizer.step() + + dist.barrier() + + if rank == 0: + print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}') + + if epoch % 5 == 0: + train_acc = test(train_loader, model, rank) + val_acc = test(val_loader, model, rank) + test_acc = test(test_loader, model, rank) + + if world_size > 1: + dist.all_reduce(train_acc, op=dist.ReduceOp.SUM) + dist.all_reduce(train_acc, op=dist.ReduceOp.SUM) + dist.all_reduce(train_acc, op=dist.ReduceOp.SUM) + train_acc /= world_size + val_acc /= world_size + test_acc /= world_size + + if rank == 0: + print(f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, ' + f'Test: {test_acc:.4f}') + + dist.barrier() + + dist.destroy_process_group() + + +if __name__ == '__main__': + dataset = Reddit('../../data/Reddit') + + world_size = torch.cuda.device_count() + print("Let's use", world_size, "GPUs!") + mp.spawn(run, args=(world_size, dataset), nprocs=world_size, join=True) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py new file mode 100644 index 0000000..2dd46cb --- /dev/null +++ b/examples/gcn_15d.py @@ -0,0 +1,1640 @@ +import argparse +import copy +import math +import os +import os.path as osp +import time +import numpy as np + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.autograd.profiler as profiler +import torch_geometric +import torch_geometric.transforms as T +from torch_geometric.data import Data, Dataset +from torch_geometric.nn import SAGEConv +from torch_geometric.datasets import Planetoid, Reddit +from torch_geometric.loader import NeighborSampler +from torch_geometric.utils import add_remaining_self_loops +import torch_sparse + +from cagnet.nn.conv import GCNConv +from cagnet.partitionings import Partitioning +from cagnet.samplers import ladies_sampler, sage_sampler +from cagnet.samplers.utils import * +import cagnet.nn.functional as CAGF +import torch.nn.functional as F + +from sparse_coo_tensor_cpp import sort_dst_proc_gpu + +import socket + +class GCN(nn.Module): + def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, partitioning, replication, + device, group=None, row_groups=None, col_groups=None): + super(GCN, self).__init__() + self.layers = nn.ModuleList() + self.n_layers = n_layers + self.aggr = aggr + self.rank = rank + self.size = size + self.group = group + self.row_groups = row_groups + self.col_groups = col_groups + self.device = device + self.partitioning = partitioning + self.replication = replication + self.timings = dict() + + self.timings["total"] = [] + self.timings["sample"] = [] + self.timings["extract"] = [] + self.timings["extract-select"] = [] + self.timings["extract-inst"] = [] + self.timings["extract-coalesce"] = [] + self.timings["train"] = [] + self.timings["selectfeats"] = [] + self.timings["fwd"] = [] + self.timings["bwd"] = [] + self.timings["loss"] = [] + self.timings["fakemats"] = [] + + self.timings["precomp"] = [] + self.timings["spmm"] = [] + self.timings["gemm_i"] = [] + self.timings["gemm_w"] = [] + self.timings["aggr"] = [] + + # # input layer + # self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # # hidden layers + # for i in range(n_layers - 2): + # self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # # output layer + # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) + # # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + + self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) + for _ in range(n_layers - 2): + self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) + self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) + + def forward(self, graphs, inputs, epoch): + h = inputs + for l, layer in enumerate(self.layers): + graphs[l].t_() + edge_index = graphs[l]._indices() + # print(f"l: {l} h.size: {h.size()} dtype: {h.dtype} graph.sizes: {graphs[l].size()} edge_index: {edge_index} dtype: {edge_index.dtype}", flush=True) + # edge_index = graphs[l] # SAGEConvfake + # edge_index, _, size = graphs[l] # quiver + h = self.layers[l](h, edge_index) # SAGEConv + # h = layer(self, graphs[l], h, epoch) # GCNConv + if l != len(self.layers) - 1: + # h = CAGF.relu(h, self.partitioning) + h = F.relu(h) + + # h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + h = F.log_softmax(h, dim=1) + return h + + @torch.no_grad() + def evaluate(self, graph, features, test_idx, labels): + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[-1], batch_size=2048, + # shuffle=False, num_workers=6) + + # for i in range(self.n_layers): + # xs = [] + # for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + # x = features[n_id].to(self.device) + # # edge_index, _, size = adj + # # x = features[n_id] + # # x_target = x[:size[1]] + # # x = self.convs[i]((x, x_target), edge_index) + # x = self.layers[i](x, edge_index) + # if i != self.n_layers - 1: + # x = F.relu(x) + # # xs.append(x) + # xs.append(x[:batch_size]) + + # features = torch.cat(xs, dim=0) + + # return features + + subgraph_loader = NeighborSampler(graph, node_idx=None, + sizes=[-1], batch_size=2048, + # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, + # sizes=[-1, -1], batch_size=512, + shuffle=False) + non_eval_timings = copy.deepcopy(self.timings) + for l, layer in enumerate(self.layers): + hs = [] + for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + edge_index, _, size = adj + adj_batch = torch.sparse_coo_tensor(edge_index, + torch.FloatTensor(edge_index.size(1)).fill_(1.0), + size) + adj_batch = adj_batch.t().coalesce() + h = features[n_id] + + # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + h = self.layers[l](h, edge_index) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + # hs.append(h) # GCNConv + hs.append(h[:batch_size]) # SAGEConv + features = torch.cat(hs, dim=0) + return features + + # print(f"test_idx: {test_idx}", flush=True) + # non_eval_timings = copy.deepcopy(self.timings) + # correct_count = 0 + # for batch_size, n_id, adjs in subgraph_loader: + # h = features[n_id] + # for l, (edge_index, _, size) in enumerate(adjs): + # # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + # h = self.layers[l](h, edge_index) # SAGEConv + # if l != len(self.layers) - 1: + # h = F.relu(h) + # h = F.log_softmax(h, dim=1) + # batch = n_id[:batch_size] + # print(f"batch: {batch}", flush=True) + # print(f"h[:batch_size]: {h[:batch_size]}", flush=True) + # print(f"h[:batch_size].argmax(dim=-1): {h[:batch_size].argmax(dim=-1)}", flush=True) + # preds = h[:batch_size].argmax(dim=-1) == labels[batch] + # correct_count += preds.sum() + # return correct_count + + # """ SAGEConv inference """ + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[-1], batch_size=2048, + # shuffle=False) + + # # Compute representations of nodes layer by layer, using *all* + # # available edges. This leads to faster computation in contrast to + # # immediately computing the final representations of each batch. + # total_edges = 0 + # for i in range(self.n_layers): + # xs = [] + # for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + # total_edges += edge_index.size(1) + # x = features[n_id].to(self.device) + # # x_target = x[:size[1]] + # # x = self.convs[i]((x, x_target), edge_index) + # x = self.layers[i](x, edge_index) # SAGEConv + # if i != self.n_layers - 1: + # x = F.relu(x) + # x = x[:batch_size] + # xs.append(x) + # features = torch.cat(xs, dim=0) + + # return features + +class LADIES(nn.Module): + def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, partitioning, replication, + device, group=None, row_groups=None, col_groups=None): + super(LADIES, self).__init__() + self.layers = nn.ModuleList() + self.n_layers = n_layers + self.aggr = aggr + self.rank = rank + self.size = size + self.group = group + self.row_groups = row_groups + self.col_groups = col_groups + self.device = device + self.partitioning = partitioning + self.replication = replication + self.timings = dict() + + self.timings["total"] = [] + self.timings["sample"] = [] + self.timings["extract"] = [] + self.timings["extract-select"] = [] + self.timings["extract-inst"] = [] + self.timings["extract-coalesce"] = [] + self.timings["train"] = [] + self.timings["selectfeats"] = [] + self.timings["fwd"] = [] + self.timings["bwd"] = [] + self.timings["loss"] = [] + self.timings["fakemats"] = [] + + self.timings["precomp"] = [] + self.timings["spmm"] = [] + self.timings["gemm_i"] = [] + self.timings["gemm_w"] = [] + self.timings["aggr"] = [] + + # # input layer + # self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # # hidden layers + # for i in range(n_layers - 2): + # self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # # output layer + # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) + self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + + def forward(self, graphs, inputs, epoch): + h = inputs + for l, layer in enumerate(self.layers): + # graphs[l].t_() + # edge_index = graphs[l]._indices() + h = layer(self, graphs[l], h, epoch) # GCNConv + if l != len(self.layers) - 1: + # h = CAGF.relu(h, self.partitioning) + h = F.relu(h) + + # h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + h = F.log_softmax(h, dim=1) + return h + + @torch.no_grad() + def evaluate(self, graph, features, test_idx, labels): + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[-1], batch_size=2048, + # shuffle=False, num_workers=6) + + # for i in range(self.n_layers): + # xs = [] + # for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + # x = features[n_id].to(self.device) + # # edge_index, _, size = adj + # # x = features[n_id] + # # x_target = x[:size[1]] + # # x = self.convs[i]((x, x_target), edge_index) + # x = self.layers[i](x, edge_index) + # if i != self.n_layers - 1: + # x = F.relu(x) + # # xs.append(x) + # xs.append(x[:batch_size]) + + # features = torch.cat(xs, dim=0) + + # return features + + subgraph_loader = NeighborSampler(graph, node_idx=None, + sizes=[-1], batch_size=2048, + # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, + # sizes=[-1, -1], batch_size=512, + shuffle=False) + non_eval_timings = copy.deepcopy(self.timings) + for l, layer in enumerate(self.layers): + hs = [] + for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + edge_index, _, size = adj + adj_batch = torch.sparse_coo_tensor(edge_index, + torch.FloatTensor(edge_index.size(1)).fill_(1.0), + size) + adj_batch = adj_batch.t().coalesce() + h = features[n_id] + + # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + h = self.layers[l](h, edge_index) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + # hs.append(h) # GCNConv + hs.append(h[:batch_size]) # SAGEConv + features = torch.cat(hs, dim=0) + return features + +def get_proc_groups(rank, size, replication): + rank_c = rank // replication + + row_procs = [] + for i in range(0, size, replication): + row_procs.append(list(range(i, i + replication))) + + col_procs = [] + for i in range(replication): + col_procs.append(list(range(i, size, replication))) + + row_groups = [] + for i in range(len(row_procs)): + row_groups.append(dist.new_group(row_procs[i])) + + col_groups = [] + for i in range(len(col_procs)): + col_groups.append(dist.new_group(col_procs[i])) + + return row_groups, col_groups + +# Normalize all elements according to KW's normalization rule +def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx, normalization): + if not normalization: + return adj_part + + adj_part = adj_part.coalesce() + deg = torch.histc(adj_matrix[0].float(), bins=node_count) + deg = deg.pow(-0.5) + + row_len = adj_part.size(0) + col_len = adj_part.size(1) + + dleft = torch.sparse_coo_tensor([np.arange(0, row_len).tolist(), + np.arange(0, row_len).tolist()], + deg[row_vtx:(row_vtx + row_len)].float(), + size=(row_len, row_len), + requires_grad=False, device=torch.device("cpu")) + + dright = torch.sparse_coo_tensor([np.arange(0, col_len).tolist(), + np.arange(0, col_len).tolist()], + deg[col_vtx:(col_vtx + col_len)].float(), + size=(col_len, col_len), + requires_grad=False, device=torch.device("cpu")) + # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) + ad_ind, ad_val = torch_sparse.spspmm(adj_part._indices(), adj_part._values(), + dright._indices(), dright._values(), + adj_part.size(0), adj_part.size(1), dright.size(1)) + + adj_part_ind, adj_part_val = torch_sparse.spspmm(dleft._indices(), dleft._values(), + ad_ind, ad_val, + dleft.size(0), dleft.size(1), adj_part.size(1)) + + adj_part = torch.sparse_coo_tensor(adj_part_ind, adj_part_val, + size=(adj_part.size(0), adj_part.size(1)), + requires_grad=False, device=torch.device("cpu")) + + return adj_part + +# Split a COO into partitions of size n_per_proc +# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. +def split_coo(adj_matrix, node_count, n_per_proc, dim): + vtx_indices = list(range(0, node_count, n_per_proc)) + vtx_indices.append(node_count) + + am_partitions = [] + for i in range(len(vtx_indices) - 1): + am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] + am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] + am_part[dim] -= vtx_indices[i] + am_partitions.append(am_part) + + return am_partitions, vtx_indices + +def row_normalize(mx): + """Row-normalize sparse matrix""" + rowsum = torch.sparse.sum(mx, 1) + r_inv = torch.float_power(rowsum, -1).flatten() + # r_inv._values = r_inv._values()[torch.isinf(r_inv._values())] = 0. + # r_mat_inv = torch.diag(r_inv._values()) + r_inv_values = torch.cuda.DoubleTensor(r_inv.size(0)).fill_(0) + r_inv_values[r_inv._indices()[0,:]] = r_inv._values() + # r_inv_values = r_inv._values() + r_inv_values[torch.isinf(r_inv_values)] = 0 + r_mat_inv = torch.sparse_coo_tensor([np.arange(0, r_inv.size(0)).tolist(), + np.arange(0, r_inv.size(0)).tolist()], + r_inv_values, + size=(r_inv.size(0), r_inv.size(0))) + # mx = r_mat_inv.mm(mx.float()) + mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values(), + mx._indices(), mx._values(), + r_mat_inv.size(0), r_mat_inv.size(1), mx.size(1), + coalesced=True) + mx = torch.sparse_coo_tensor(indices=mx_indices, values=mx_values.double(), size=(r_mat_inv.size(0), mx.size(1))) + return mx + +def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, \ + normalize, replicate_graph): + node_count = inputs.size(0) + # n_per_proc = math.ceil(float(node_count) / size) + # n_per_proc = math.ceil(float(node_count) / (size / replication)) + n_per_proc = node_count // (size // replication) + + am_partitions = None + am_pbyp = None + + # inputs = inputs.to(torch.device("cpu")) + # adj_matrix = adj_matrix.to(torch.device("cpu")) + # torch.cuda.synchronize() + + rank_c = rank // replication + # Compute the adj_matrix and inputs partitions for this process + # TODO: Maybe I do want grad here. Unsure. + with torch.no_grad(): + # Column partitions + if not replicate_graph: + am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) + else: + am_partitions = None + # # proc_node_count = vtx_indices[rank_c + 1] - vtx_indices[rank_c] + # # am_pbyp, _ = split_coo(am_partitions[rank_c], node_count, n_per_proc, 0) + # # print(f"before", flush=True) + # # for i in range(len(am_pbyp)): + # # if i == size // replication - 1: + # # last_node_count = vtx_indices[i + 1] - vtx_indices[i] + # # am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), + # # size=(last_node_count, proc_node_count), + # # requires_grad=False) + + # # am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + # # vtx_indices[rank_c], normalize) + # # else: + # # am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), + # # size=(n_per_proc, proc_node_count), + # # requires_grad=False) + + # # am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + # # vtx_indices[rank_c], normalize) + + if not replicate_graph: + for i in range(len(am_partitions)): + proc_node_count = vtx_indices[i + 1] - vtx_indices[i] + am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], + torch.ones(am_partitions[i].size(1)).double(), + size=(node_count, proc_node_count), + requires_grad=False) + am_partitions[i] = scale_elements(adj_matrix, am_partitions[i], node_count, 0, vtx_indices[i], \ + normalize) + adj_matrix_loc = am_partitions[rank_c] + else: + adj_matrix_loc = None + + # # input_partitions = torch.split(inputs, math.ceil(float(inputs.size(0)) / (size / replication)), dim=0) + input_partitions = torch.split(inputs, inputs.size(0) // (size // replication), dim=0) + if len(input_partitions) > (size // replication): + input_partitions_fused = [None] * (size // replication) + input_partitions_fused[:-1] = input_partitions[:-2] + input_partitions_fused[-1] = torch.cat(input_partitions[-2:], dim=0) + input_partitions = input_partitions_fused + + inputs_loc = input_partitions[rank_c] + + # print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) + print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}", flush=True) + return inputs_loc, adj_matrix_loc, am_partitions, input_partitions + +def one5d_partition_mb(rank, size, batches, replication, mb_count): + rank_c = rank // replication + batch_partitions = torch.split(batches, int(mb_count // (size / replication)), dim=0) + return batch_partitions[rank_c] + # batch_partitions = torch.split(batches, int(mb_count // size), dim=0) + # return batch_partitions[rank] + +def main(args, batches=None): + # load and preprocess dataset + # Initialize distributed environment with SLURM + if "SLURM_PROCID" in os.environ.keys(): + os.environ["RANK"] = os.environ["SLURM_PROCID"] + + if "SLURM_NTASKS" in os.environ.keys(): + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + + os.environ["MASTER_ADDR"] = args.hostname + os.environ["MASTER_PORT"] = "1234" + + print(f"device_count: {torch.cuda.device_count()}") + print(f"hostname: {socket.gethostname()}", flush=True) + if not dist.is_initialized(): + dist.init_process_group(backend=args.dist_backend) + rank = dist.get_rank() + size = dist.get_world_size() + print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) + torch.cuda.set_device(rank % args.gpu) + + device = torch.device(f'cuda:{rank % args.gpu}') + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + start_inner_timer = torch.cuda.Event(enable_timing=True) + stop_inner_timer = torch.cuda.Event(enable_timing=True) + + path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', args.dataset) + + if args.dataset == "cora" or args.dataset == "reddit": + if args.dataset == "cora": + dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures()) + elif args.dataset == "reddit": + dataset = Reddit(path) + + data = dataset[0] + data = data.to(device) + # data.x.requires_grad = True + inputs = data.x.to(device) + # inputs.requires_grad = True + data.y = data.y.to(device) + edge_index = data.edge_index + num_features = dataset.num_features + num_classes = dataset.num_classes + adj_matrix = edge_index + elif args.dataset == "Amazon": + print(f"Loading coo...", flush=True) + edge_index = torch.load("/global/homes/a/alokt/data/Amazon/processed/data.pt") + print(f"Done loading coo", flush=True) + # n = 14249639 + n = 9430088 + num_features = 128 + num_classes = 24 + inputs = torch.rand(n, num_features) + data = Data() + data.y = torch.rand(n).uniform_(0, num_classes - 1).long() + # data.train_mask = torch.ones(n).long() + data.train_mask = torch.zeros(n).long() + data.train_mask[:args.batch_size * 1024] = 1 + data.test_mask = torch.zeros(n).long() + data.test_mask[args.batch_size * args.n_bulkmb:] = 1 + adj_matrix = edge_index.t_() + data = data.to(device) + # inputs.requires_grad = True + data.y = data.y.to(device) + elif args.dataset == "Protein": + print(f"Loading coo...", flush=True) + edge_index = torch.load("../../data/protein/processed/protein.pt") + print(f"Done loading coo", flush=True) + n = 8745542 + num_features = 128 + num_classes = 256 + inputs = torch.rand(n, num_features) + data = Data() + data.y = torch.rand(n).uniform_(0, num_classes - 1).long() + # data.train_mask = torch.ones(n).long() + data.train_mask = torch.zeros(n).long() + data.train_mask[:args.batch_size * 1024] = 1 + data.test_mask = torch.ones(n).long() + adj_matrix = edge_index.t_() + data = data.to(device) + # inputs.requires_grad = True + data.y = data.y.to(device) + elif args.dataset == "Protein_sg2": + n = 17491085 + num_features = 128 + num_classes = 256 + if rank % args.gpu == 0: + print(f"Loading coo...", flush=True) + edge_index = torch.load("../../data/protein_sg2/processed/protein_sg2.pt") + print(f"Done loading coo", flush=True) + inputs = torch.rand(n, num_features) + adj_matrix = edge_index.t_() + # inputs.requires_grad = True + data = Data() + data.y = torch.rand(n).uniform_(0, num_classes - 1).long() + data.train_mask = torch.ones(n).long() + data.test_mask = torch.ones(n).long() + data = data.to(device) + data.y = data.y.to(device) + elif args.dataset == "ogbn-products": + import ogb + data = Data() + from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time + dataset = PygNodePropPredDataset(name=args.dataset, root="/global/u1/a/alokt/data") + data = dataset[0] + + # split_idx = dataset.get_idx_split() + # # train_loader = DataLoader(dataset[split_idx["train"]], batch_size=32, shuffle=True) + # # valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=32, shuffle=False) + # # test_loader = DataLoader(dataset[split_idx["test"]], batch_size=32, shuffle=False) + # # data = data.to(device) + # # data.x.requires_grad = True + # # inputs = data.x.to(device) + inputs = data.x + data.y = data.y.squeeze().to(device) + # # inputs.requires_grad = True + # # data.y = data.y.to(device) + num_features = dataset.num_features + edge_index = data.edge_index + num_classes = dataset.num_classes + adj_matrix = edge_index + split_idx = dataset.get_idx_split() + train_idx = split_idx['train'].to(device) + test_idx = split_idx['test'].to(device) + elif args.dataset == "ogbn-papers100M": + # import ogb + data = Data() + # from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time + if ("papers100M" in args.dataset and rank % args.gpu == 0) or "products" in args.dataset: + # dataset = PygNodePropPredDataset(name=args.dataset, root="/global/u1/a/alokt/data") + # data = dataset[0] + + # # split_idx = dataset.get_idx_split() + # # # train_loader = DataLoader(dataset[split_idx["train"]], batch_size=32, shuffle=True) + # # # valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=32, shuffle=False) + # # # test_loader = DataLoader(dataset[split_idx["test"]], batch_size=32, shuffle=False) + # # # data = data.to(device) + # # # data.x.requires_grad = True + # # # inputs = data.x.to(device) + # inputs = data.x + # # data.y = data.y.squeeze().to(device) + # # # inputs.requires_grad = True + # # # data.y = data.y.to(device) + # num_features = dataset.num_features + # # edge_index = data.edge_index + # num_classes = dataset.num_classes + # adj_matrix = edge_index + # split_idx = dataset.get_idx_split() + # train_idx = split_idx['train'].to(device) + # test_idx = split_idx['test'].to(device) + if not args.replicate_graph: + import ogb + from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time + dataset = PygNodePropPredDataset(name=args.dataset, root="/global/u1/a/alokt/data") + data = dataset[0] + + # split_idx = dataset.get_idx_split() + # # train_loader = DataLoader(dataset[split_idx["train"]], batch_size=32, shuffle=True) + # # valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=32, shuffle=False) + # # test_loader = DataLoader(dataset[split_idx["test"]], batch_size=32, shuffle=False) + # # data = data.to(device) + # # data.x.requires_grad = True + # # inputs = data.x.to(device) + inputs = data.x + data.y = data.y.squeeze().to(device) + # # inputs.requires_grad = True + # # data.y = data.y.to(device) + num_features = dataset.num_features + edge_index = data.edge_index + num_classes = dataset.num_classes + adj_matrix = edge_index + split_idx = dataset.get_idx_split() + train_idx = split_idx['train'].to(device) + test_idx = split_idx['test'].to(device) + else: + adj_matrix = None + print(f"before loading", flush=True) + data.y = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/label/label.pt") + data.y = data.y.squeeze().to(device) + print(f"data.y: {data.y} data.y.dtype {data.y.dtype}", flush=True) + train_idx = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/index/train_idx.pt") + train_idx = train_idx.to(device) + print(f"train_idx: {train_idx} train_idx.dtype {train_idx.dtype}", flush=True) + inputs = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") + print(f"inputs: {inputs} inputs.dtype {inputs.dtype}", flush=True) + num_features = inputs.size(1) + num_classes = 172 + + # print(f"adj_matrix.size(): {adj_matrix.size()}") + # print(f"adj_matrix: {adj_matrix}") + + # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + # if args.sample_method == "ladies": + # print(f"before adj_matrix.size: {adj_matrix.size()}", flush=True) + # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + # print(f"after adj_matrix.size: {adj_matrix.size()}", flush=True) + if args.dataset != "ogbn-papers100M" and args.dataset != "Protein_sg2": + edge_count = adj_matrix.size(1) + + partitioning = Partitioning.ONE5D + + torch.cuda.synchronize() + row_groups, col_groups = get_proc_groups(rank, size, args.replication) + + proc_row = size // args.replication + rank_row = rank // args.replication + rank_col = rank % args.replication + if rank_row >= (size // args.replication): + return + + torch.cuda.synchronize() + if args.dataset == "ogbn-papers100M" or "Protein" in args.dataset: + if rank % args.gpu == 0: + inputs = inputs.to(torch.device("cpu")) + if args.replicate_graph: + # send g_loc.nnz, g_loc, train_idx.len, and train_idx + features_loc, g_loc, am_partitions, input_partitions = \ + one5d_partition(rank, size, inputs, adj_matrix, data, + inputs, num_classes, args.replication, \ + args.normalize, args.replicate_graph) + + if args.dataset == "ogbn-papers100M": + print(f"begin copying", flush=True) + adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") + adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") + print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) + print(f"adj_matrix_cols: {adj_matrix_cols} dtype: {adj_matrix_cols.dtype} device: {adj_matrix_cols.device}", flush=True) + # adj_matrix_crows = adj_matrix_crows.int().to(device) + # adj_matrix_cols = adj_matrix[1,:].int().to(device) + adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) + g_loc = torch.sparse_csr_tensor(adj_matrix_crows.int(), adj_matrix_cols.int(), \ + adj_matrix_vals, \ + torch.Size([inputs.size(0), inputs.size(0)])) + else: + print(f"begin construction", flush=True) + adj_matrix_vals = torch.FloatTensor(adj_matrix.size(1)).fill_(1) + g_loc = torch.sparse_coo_tensor(adj_matrix, adj_matrix_vals, \ + torch.Size([inputs.size(0), inputs.size(0)])) + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, adj_matrix_vals, g_loc.size()) + g_loc = g_loc.to(device) + else: + adj_matrix = adj_matrix.to(torch.device("cpu")) + features_loc, g_loc, am_partitions, input_partitions = \ + one5d_partition(rank, size, inputs, adj_matrix, data, + inputs, num_classes, args.replication, \ + args.normalize, args.replicate_graph) + features_zero_row = input_partitions[0][0].to(device) # feature vector in row 0 + print(f"stop copying", flush=True) + for i in range(1, args.gpu): + print(f"iteration i: {i}", flush=True) + if args.replicate_graph: + rank_send_row = (rank + i) // args.replication + dst_gpu = rank + i + g_send = adj_matrix + print("normalizing", flush=True) + + features_send = input_partitions[rank_send_row].to(device) + edge_count = g_loc.col_indices().size(0) + + + g_send_meta = torch.cuda.LongTensor(6) + g_send_meta[0] = g_loc.col_indices().size(0) + g_send_meta[1] = inputs.size(0) + g_send_meta[2] = inputs.size(0) + g_send_meta[3] = g_loc.col_indices().size(0) + g_send_meta[4] = num_features + g_send_meta[5] = num_classes + dist.send(g_send_meta, dst=dst_gpu) + dist.send(g_loc.crow_indices(), dst=dst_gpu) + dist.send(g_loc.col_indices(), dst=dst_gpu) + dist.send(features_send, dst=dst_gpu) + dist.send(features_zero_row, dst=dst_gpu) + dist.send(data.y.float(), dst=dst_gpu) + else: + rank_send_row = (rank + i) // args.replication + dst_gpu = rank + i + g_send = am_partitions[rank_send_row] + print("coalescing", flush=True) + g_send = g_send.coalesce().t_() + print("normalizing", flush=True) + g_send = g_send.to(device) + g_send = g_send.double() + features_send = input_partitions[rank_send_row].to(device) + edge_count = adj_matrix.size(1) + + + g_send_meta = torch.cuda.LongTensor(6) + g_send_meta[0] = g_send._nnz() + g_send_meta[1] = g_send.size(0) + g_send_meta[2] = g_send.size(1) + g_send_meta[3] = adj_matrix.size(1) + g_send_meta[4] = num_features + g_send_meta[5] = num_classes + dist.send(g_send_meta, dst=dst_gpu) + dist.send(g_send._indices(), dst=dst_gpu) + dist.send(g_send._values(), dst=dst_gpu) + dist.send(features_send, dst=dst_gpu) + dist.send(features_zero_row, dst=dst_gpu) + dist.send(data.y.float(), dst=dst_gpu) + + del g_send + del features_send + + if args.dataset == "ogbn-papers100M": + train_idx_len = torch.cuda.LongTensor(1).fill_(train_idx.size(0)) + # test_idx_len = torch.cuda.LongTensor(1).fill_(test_idx.size(0)) + dist.send(train_idx_len, dst=dst_gpu) + # dist.send(test_idx_len, dst=dst_gpu) + dist.send(train_idx, dst=dst_gpu) + # dist.send(test_idx, dst=dst_gpu) + if args.sample_method == "sage": + if args.replicate_graph: + edge_count = g_loc.col_indices().size(0) + else: + edge_count = g_loc._indices().size(1) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t_() + g_loc = g_loc.cpu() + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), g_loc.size()) + print("host2device", flush=True) + g_loc = g_loc.to(device) + else: + g_loc = g_loc.coalesce().t_() + g_loc = g_loc.to(device) + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), g_loc.size()) + features_loc = features_loc.to(device) + + # features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + # args.replication, args.normalize) + # g_send_meta = torch.cuda.LongTensor(6) + # g_send_meta[0] = g_loc._nnz() + # g_send_meta[1] = g_loc.size(0) + # g_send_meta[2] = g_loc.size(1) + # g_send_meta[3] = adj_matrix.size(1) + # g_send_meta[4] = num_features + # g_send_meta[5] = num_classes + # torch.save(g_send_meta, f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_indices.pt") + # torch.save(g_loc._indices(), f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_indices.pt") + # torch.save(g_loc._values(), f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_values.pt") + # torch.save(features_loc, f"../../data/ogbn_papers100M/cagnet/p16/{rank}features_loc.pt") + + # del adj_matrix # Comment when testing + # del inputs + else: + src_gpu = rank - (rank % args.gpu) + if args.replicate_graph: + g_loc_meta = torch.cuda.LongTensor(6).fill_(0) + dist.recv(g_loc_meta, src=src_gpu) + + # g_loc_indices = torch.cuda.IntTensor(2, g_loc_meta[0].item()) + # dist.recv(g_loc_indices, src=src_gpu) + # torch.cuda.synchronize() + # print(f"recv indices", flush=True) + + g_loc_crows = torch.cuda.IntTensor(g_loc_meta[1].item() + 1) + dist.recv(g_loc_crows, src=src_gpu) + torch.cuda.synchronize() + print(f"recv crows", flush=True) + + g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) + dist.recv(g_loc_cols, src=src_gpu) + torch.cuda.synchronize() + print(f"recv cols", flush=True) + + # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + # dist.recv(g_loc_values, src=src_gpu) + + num_features = g_loc_meta[4].item() + num_classes = g_loc_meta[5].item() + rows_per_proc = g_loc_meta[1].item() // (size // args.replication) + if rank_row < proc_row - 1: + inputs = torch.cuda.FloatTensor(rows_per_proc, num_features) + else: + rows = g_loc_meta[1].item() - rows_per_proc * (proc_row - 1) + inputs = torch.cuda.FloatTensor(rows, num_features) + dist.recv(inputs, src=src_gpu) + torch.cuda.synchronize() + print(f"recv inputs", flush=True) + + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + print(f"data.y.size: {data.y.size()}", flush=True) + dist.recv(data.y, src=src_gpu) + torch.cuda.synchronize() + data.y = data.y.long() + + # g_loc_values = torch.cuda.FloatTensor(g_loc_indices.size(1)).fill_(1.0) + # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + # torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc_values = torch.cuda.FloatTensor(g_loc_cols.size(0)).fill_(1.0) + g_loc_crows = g_loc_crows.cpu() + g_loc_cols = g_loc_cols.cpu() + g_loc_values = g_loc_values.cpu() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc = g_loc.to(device) + edge_count = g_loc_meta[3].item() + else: + src_gpu = rank - (rank % args.gpu) + g_loc_meta = torch.cuda.LongTensor(6).fill_(0) + dist.recv(g_loc_meta, src=src_gpu) + + g_loc_indices = torch.cuda.LongTensor(2, g_loc_meta[0].item()) + dist.recv(g_loc_indices, src=src_gpu) + + g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + dist.recv(g_loc_values, src=src_gpu) + + num_features = g_loc_meta[4].item() + num_classes = g_loc_meta[5].item() + if rank_row < proc_row - 1: + inputs = torch.cuda.FloatTensor(g_loc_meta[1].item(), num_features) + else: + rows = g_loc_meta[1].item() + (g_loc_meta[2].item() - proc_row * g_loc_meta[1].item()) + inputs = torch.cuda.FloatTensor(rows, num_features) + dist.recv(inputs, src=src_gpu) + + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + dist.recv(data.y, src=src_gpu) + data.y = data.y.long() + + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + if args.sample_method == "sage": + g_loc = g_loc.cpu() + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), g_loc.size()) + g_loc = g_loc.to(device) + else: + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), g_loc.size()) + + edge_count = g_loc_meta[3].item() + + if args.dataset == "ogbn-papers100M": + train_idx_len = torch.cuda.LongTensor(1).fill_(0) + # test_idx_len = torch.cuda.LongTensor(1).fill_(0) + dist.recv(train_idx_len, src=src_gpu) + # dist.recv(test_idx_len, src=src_gpu) + + train_idx = torch.cuda.LongTensor(train_idx_len.item()).fill_(0) + # test_idx = torch.cuda.LongTensor(test_idx_len.item()).fill_(0) + dist.recv(train_idx, src=src_gpu) + # dist.recv(test_idx, src=src_gpu) + torch.cuda.synchronize() + features_loc = inputs.to(device) + else: + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) + features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + args.replication, args.normalize, args.replicate_graph) + print("coalescing", flush=True) + if args.dataset == "ogbn-products" or args.dataset == "Amazon": + adj_matrix_values = torch.FloatTensor(adj_matrix.size(1)).fill_(1.0) + g_loc = torch.sparse_coo_tensor(adj_matrix.int(), adj_matrix_values, + size=torch.Size([inputs.size(0), inputs.size(0)])) + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + adj_matrix_values = torch.FloatTensor(g_loc_cols.size(0)).fill_(1.0) + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, adj_matrix_values, + size=torch.Size([inputs.size(0), inputs.size(0)])) + else: + if args.replicate_graph: + adj_matrix_row_lens = torch.IntTensor(inputs.size(0)).fill_(0) + ones = torch.IntTensor(adj_matrix.size(1)).fill_(1) + adj_matrix_row_lens.scatter_add_(0, adj_matrix[0,:].long(), ones) + adj_matrix_crows = torch.IntTensor(inputs.size(0) + 1) + adj_matrix_crows[1:] = torch.cumsum(adj_matrix_row_lens, dim=0) + adj_matrix_crows[0] = 0 + adj_matrix_values = torch.FloatTensor(adj_matrix.size(1)).fill_(1.0) + g_loc = torch.sparse_csr_tensor(adj_matrix_crows, adj_matrix[1,:].int(), adj_matrix_values, + size=torch.Size([inputs.size(0), inputs.size(0)])) + else: + g_loc = g_loc.coalesce().t_() + g_loc = g_loc.to_sparse_csr() + g_loc_crows = g_loc.crow_indices().int() + g_loc_cols = g_loc.col_indices().int() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), g_loc.size()) + print("normalizing", flush=True) + g_loc = g_loc.to(device) + # g_loc = g_loc.double() + # features_loc = inputs + features_zero_row = inputs[0].to(device) + features_loc = features_loc.to(device) + # del inputs + + print("end partitioning", flush=True) + print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) + print(f"features_loc.size: {features_loc.size()}", flush=True) + dist.barrier() + + if rank == 0 and args.dataset != "ogbn-papers100M": + adj_matrix = adj_matrix.cpu() + inputs = inputs.cpu() + + # g_loc_indices, _ = add_remaining_self_loops(g_loc._indices(), num_nodes=g_loc.size(0)) + # g_loc_values = torch.cuda.DoubleTensor(g_loc_indices.size(1)).fill_(1) + # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, g_loc.size()) + # g_loc = row_normalize(g_loc) + print("done normalizing", flush=True) + torch.set_printoptions(precision=10) + + # n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) + n_per_proc = g_loc.size(0) // (size // args.replication) + + if not args.dataset.startswith("ogb"): + train_nid = data.train_mask.nonzero().squeeze() + test_nid = data.test_mask.nonzero().squeeze() + print(f"train_nid.size: {train_nid.size()}") + print(f"test_nid.size: {test_nid.size()}") + else: + if args.dataset == "ogbn-papers100M": + train_nid = train_idx + print(f"train_nid.size: {train_nid.size()}") + else: + train_nid = train_idx + test_nid = test_idx + print(f"train_nid.size: {train_nid.size()}") + print(f"test_nid.size: {test_nid.size()}") + + # # # adj_matrix = None # Uncomment bottom for testing + # adj_matrix = adj_matrix.cuda() + # adj_matrix = torch.sparse_coo_tensor(adj_matrix, + # torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0), + # size=(node_count, node_count)) + # adj_matrix = adj_matrix.coalesce().cuda().double() + + # return frontiers, adj_matrices, adj_matrix + # return frontiers, adj_matrices, adj_matrix, col_groups + + # create GCN model + # torch.manual_seed(0) + if args.sample_method == "sage": + model = GCN(num_features, + args.n_hidden, + num_classes, + args.n_layers, + args.aggr, + rank, + size, + Partitioning.NONE, + args.replication, + device, + row_groups=row_groups, + col_groups=col_groups) + elif args.sample_method == "ladies": + model = LADIES(num_features, + args.n_hidden, + num_classes, + args.n_layers, + args.aggr, + rank, + size, + Partitioning.NONE, + args.replication, + device, + row_groups=row_groups, + col_groups=col_groups) + model = model.to(device) + + # use optimizer + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) + # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + + dur = [] + # torch.manual_seed(0) + # tally = torch.cuda.IntTensor(size).fill_(0) + tally = torch.cuda.IntTensor(proc_row).fill_(0) + row_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + if rank == size - 1: + row_n_bulkmb = args.n_bulkmb - row_n_bulkmb * (proc_row - 1) + + rank_n_bulkmb = int(row_n_bulkmb / args.replication) + if rank == size - 1: + rank_n_bulkmb = row_n_bulkmb - rank_n_bulkmb * (args.replication - 1) + + print(f"rank_n_bulkmb: {rank_n_bulkmb}") + # g_loc = g_loc.cpu() + + # g_loc = g_loc.to_sparse_csr() + # g_loc_crows = g_loc.crow_indices().int() + # g_loc_cols = g_loc.col_indices().int() + # g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), \ + # g_loc.size()) + g_loc = g_loc.to(device) + # csr_topo = quiver.CSRTopo(adj_matrix) + # quiver_sampler = quiver.pyg.GraphSageSampler(csr_topo, [args.samp_num] * args.n_layers, 0, + # mode='UVA') + # train_loader = torch.utils.data.DataLoader(train_nid, batch_size=args.batch_size, + # shuffle=True, drop_last=True) + for epoch in range(args.n_epochs): + # train_loader_iter = iter(train_loader) + print(f"Epoch: {epoch}", flush=True) + if epoch >= 1: + epoch_start = time.time() + model.train() + + feat_dims = [] + adj0_dims = [] + adj1_dims = [] + adj0_nnzs = [] + adj1_nnzs = [] + adj_nnzs = [] + + # print("Constructing batches", flush=True) + torch.cuda.nvtx.range_push("nvtx-construct-batches") + batch_count = -(train_nid.size(0) // -args.batch_size) # ceil(train_nid.size(0) / batch_size) + # if batches is None: + torch.manual_seed(epoch) + vertex_perm = torch.randperm(train_nid.size(0)) + batches_all = train_nid[vertex_perm].int() + torch.cuda.nvtx.range_pop() # construct-batches + + # torch.cuda.nvtx.range_push("nvtx-construct-batches") + # batch_count = -(train_nid.size(0) // -args.batch_size) # ceil(train_nid.size(0) / batch_size) + # # print(f"batch_count: {batch_count}", flush=True) + # if batches is None: + # torch.manual_seed(epoch) + # vertex_perm = torch.randperm(train_nid.size(0)) + # batches_all = train_nid[vertex_perm] + # torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("nvtx-bulk") + last_batch = False + print(f"batch_count: {batch_count}", flush=True) + for b in range(0, batch_count, args.n_bulkmb): + if b + args.n_bulkmb > batch_count: + break + last_batch = True + tmp_bulkmb = args.n_bulkmb + tmp_batch_size = args.batch_size + args.n_bulkmb = 1 + args.batch_size = batch_count - b + torch.cuda.nvtx.range_push("nvtx-partition-batches") + batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) + if args.sample_method == "sage": + batches_loc = one5d_partition_mb(rank, size, batches, 1, args.n_bulkmb) + elif args.sample_method == "ladies": + batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) + + batches_indices_rows = torch.arange(batches_loc.size(0), dtype=torch.int32, device=device) + batches_indices_rows = batches_indices_rows.repeat_interleave(batches_loc.size(1)) + batches_indices_cols = batches_loc.view(-1) + batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) + batches_values = torch.cuda.FloatTensor(batches_loc.size(1) * batches_loc.size(0), device=device).fill_(1.0) + batches_loc = torch.sparse_coo_tensor(batches_indices, batches_values, (batches_loc.size(0), g_loc.size(1))) + torch.cuda.nvtx.range_pop() # partition-batches + # g_loc = torch.pow(g_loc, 2) + + node_count = g_loc.size(1) + # adj_matrix = adj_matrix.cuda() + if args.n_darts == -1: + avg_degree = int(edge_count / node_count) + if args.sample_method == "ladies": + args.n_darts = avg_degree * args.batch_size * (args.replication ** 2) + elif args.sample_method == "sage": + args.n_darts = [] + avg_degree = edge_count / node_count + for d in range(args.n_layers): + dart_count = int(avg_degree * args.samp_num[d] / avg_degree) + args.n_darts.append(dart_count) + + # for sa-spgemm + # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) + nnz_row_masks = None + if False and not args.replicate_graph: + nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._nnz()) + nnz_row_masks.fill_(0) + + # nnz_recv_upperbound = edge_count // (size // args.replication) + + print("Sampling", flush=True) + if epoch >= 1: + start_time(start_timer, timing_arg=True) + torch.cuda.nvtx.range_push("nvtx-sampling") + if args.sample_method == "ladies": + current_frontier, next_frontier, adj_matrices_bulk = \ + ladies_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.semibulk, args.replication, \ + nnz_row_masks, rank, size, + row_groups, col_groups, args.timing, + args.replicate_graph) + frontiers_bulk = [current_frontier, next_frontier] + elif args.sample_method == "sage": + if nnz_row_masks is not None: + nnz_row_masks.fill_(False) + if args.replicate_graph: + rep_pass = 1 + else: + rep_pass = args.replication + frontiers_bulk, adj_matrices_bulk = sage_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + rep_pass, nnz_row_masks, + rank, size, row_groups, + col_groups, args.timing, args.baseline, + args.replicate_graph) + + if epoch >= 1: + model.timings["sample"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + torch.cuda.nvtx.range_pop() # nvtx-sampling + + if epoch >= 1: + start_time(start_timer, timing_arg=True) + # g_loc = g_loc.coalesce() + + print("Training", flush=True) + torch.cuda.nvtx.range_push("nvtx-training") + + for i in range(rank_n_bulkmb): + # forward + # batch_vtxs = frontiers[0][i].view(-1) + torch.cuda.nvtx.range_push("nvtx-extracting") + + if args.sample_method == "sage": + batch_select_size = args.batch_size + batch_select_min = i * args.batch_size + batch_select_max = (i + 1) * args.batch_size + + batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) + + src_select_min = i * args.batch_size * np.prod(args.samp_num[:-1], dtype=int) + src_select_max = (i + 1) * args.batch_size * np.prod(args.samp_num[:-1], dtype=int) + src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) + elif args.sample_method == "ladies": + batch_vtxs = frontiers_bulk[0][i] + src_vtxs = frontiers_bulk[-1][i] + + adjs = [None] * args.n_layers + adj_sample_skip_cols = None + for l in range(args.n_layers): + # adj_row_select_min = i * args.batch_size * (args.samp_num ** l) + # adj_row_select_max = (i + 1) * args.batch_size * (args.samp_num ** l) + adj_row_select_min = i * args.batch_size * np.prod(args.samp_num[:l], dtype=int) + adj_row_select_max = (i + 1) * args.batch_size * np.prod(args.samp_num[:l], dtype=int) + adj_row_select_max = min(adj_row_select_max, adj_matrices_bulk[l].size(0)) + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + + adj_nnz_start = adj_matrices_bulk[l].crow_indices()[adj_row_select_min] + adj_nnz_stop = adj_matrices_bulk[l].crow_indices()[adj_row_select_max] + + adj_sample_crows = adj_matrices_bulk[l].crow_indices()[adj_row_select_min:(adj_row_select_max+1)].clone() + adj_sample_cols = adj_matrices_bulk[l].col_indices()[adj_nnz_start:adj_nnz_stop] + crows_start = adj_sample_crows[0].item() + adj_sample_crows -= crows_start + + adj_sample_values = adj_matrices_bulk[l].values()[adj_nnz_start:adj_nnz_stop].float() + + if epoch >= 1: + model.timings["extract-select"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + if args.sample_method == "ladies": + adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + row_count = adj_row_select_max - adj_row_select_min + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, row_count, + device=adj_sample_row_lens.device), + adj_sample_row_lens) + adj_sample_indices = torch.stack((adj_sample_rows, adj_sample_cols)) + # adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ + # adj_sample_values, \ + # # (args.batch_size * (args.samp_num ** l), \ + # # args.batch_size * (args.samp_num ** (l + 1)))) + # (args.batch_size * np.prod(args.samp_num[:l], dtype=int), \ + # args.batch_size * (np.prod(args.samp_num[:(l+1)], dtype=int)))) + adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ + adj_sample_values, \ + (args.batch_size, args.samp_num[0] + args.batch_size)) + # (args.batch_size, args.samp_num + args.batch_size)) + else: + # adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + # row_count = adj_row_select_max - adj_row_select_min + # adj_sample_rows = torch.repeat_interleave( + # torch.arange(0, row_count, + # device=adj_sample_row_lens.device), + # adj_sample_row_lens) + # adj_sample_indices = torch.stack((adj_sample_rows, adj_sample_cols)) + # adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ + # adj_sample_values, \ + # # (args.batch_size * (args.samp_num ** l), \ + # # args.batch_size * (args.samp_num ** (l + 1)))) + # (args.batch_size * np.prod(args.samp_num[:l], dtype=int), \ + # args.batch_size * (np.prod(args.samp_num[:(l+1)], dtype=int)))) + # adj_nnzs.append(adj_sample_values.size(0)) + if l > 0: + adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + # adj_sample_row_lens[adj_sample_skip_cols] = 0 + adj_sample_row_lens[adj_sample_skip_cols] = -1 + row_mask = adj_sample_row_lens > -1 + adj_sample_row_lens = adj_sample_row_lens[row_mask] + row_count = adj_sample_row_lens.size(0) + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, row_count, + device=adj_sample_row_lens.device), + adj_sample_row_lens) + nnz_col_mask = torch.cuda.BoolTensor(adj_sample_crows[-1].item()).fill_(False) + rowselect_csr_gpu(adj_sample_skip_cols, adj_sample_crows, nnz_col_mask, \ + adj_sample_skip_cols.size(0), adj_sample_crows[-1]) + adj_sample_values = adj_sample_values[~nnz_col_mask] + adj_sample_cols = adj_sample_cols[~nnz_col_mask] + else: + adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + row_count = adj_row_select_max - adj_row_select_min + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, row_count, + device=adj_sample_row_lens.device), + adj_sample_row_lens) + + # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) + # row_count = args.batch_size * np.prod(args.samp_num[:l], dtype=int) + col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) + adj_sample_skip_cols = torch.histc(adj_sample_cols, bins=col_count) + adj_sample_skip_cols = (adj_sample_skip_cols == 0).nonzero().squeeze(1) + # print(f"l: {l} adj_sample_cols: {adj_sample_cols} adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) - \ + adj_sample_skip_cols.size(0) + adj_sample_cols = torch.arange(0, adj_sample_values.size(0), device=device, dtype=int) + # print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + # print(f"adj_sample_cols: {adj_sample_cols}", flush=True) + adj_sample_indices = torch.stack((adj_sample_rows, adj_sample_cols)) + adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ + adj_sample_values, (row_count, col_count)) + # (args.batch_size * (args.samp_num ** l), \ + # args.batch_size * (args.samp_num ** (l + 1)))) + adj_nnzs.append(adj_sample_values.size(0)) + # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") + if epoch >= 1: + model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + # adjs[args.n_layers - l - 1] = adj_matrix_sample.to_sparse_coo() + adjs[args.n_layers - l - 1] = adj_matrix_sample.coalesce() + # print(f"l: {l} size: {adjs[args.n_layers - l - 1].size()} nnz: {adjs[args.n_layers - l - 1]._nnz()}") + # print(f"l: {l} {adjs[args.n_layers - l- 1]}") + if epoch >= 1: + model.timings["extract-coalesce"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + torch.cuda.nvtx.range_pop() # nvtx-extracting + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("nvtx-selectfeats") + if size > 1: + + src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) + og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) + # src_vtxs_nnz = src_vtxs[src_vtxs_nnz_mask] + # src_vtxs_sort = torch.cuda.LongTensor(src_vtxs_nnz.size(0)) + # og_idxs = torch.cuda.LongTensor(src_vtxs_nnz.size(0)) + tally.fill_(0) + + # node_per_row = int(math.ceil(node_count / (size / args.replication))) + # node_per_proc = int(math.ceil(node_count / size)) + node_per_row = node_count // (size // args.replication) + node_per_proc = node_count // size + # sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_proc) + sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, + node_per_row, proc_row) + + src_vtxs_sort_nnz = src_vtxs_sort != 0 + tally[0] -= (src_vtxs_sort == 0).sum() + src_vtxs_sort = src_vtxs_sort[src_vtxs_sort_nnz] + og_idxs_nnz = og_idxs[src_vtxs_sort_nnz] + src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) + + output_tally = [] + for j in range(proc_row): + output_tally.append(torch.cuda.IntTensor(1)) + input_tally = list(torch.split(tally, 1)) + + dist.all_to_all(output_tally, input_tally, col_groups[rank_col]) + + output_src_vtxs = torch.cuda.LongTensor(sum(output_tally).item()) + dist.all_to_all_single(output_src_vtxs, src_vtxs_sort, + output_tally, input_tally, + col_groups[rank_col]) + output_src_vtxs -= node_per_row * rank_row + # output_src_vtxs -= node_per_proc * rank + + input_features = features_loc[output_src_vtxs] + # output_features = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) + output_features = torch.cuda.FloatTensor(src_vtxs_sort.size(0), features_loc.size(1)) + dist.all_to_all_single(output_features, input_features, + input_tally, output_tally, + col_groups[rank_col]) + + #features_batch = torch.cuda.FloatTensor(output_features.size()) + features_batch = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) + features_batch[og_idxs_nnz] = output_features + features_batch[og_idxs[~src_vtxs_sort_nnz]] = features_zero_row + if adj_sample_skip_cols is not None: + features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + features_mask[adj_sample_skip_cols] = False + features_batch = features_batch[features_mask] + else: + features_batch = features_loc[src_vtxs] + if adj_sample_skip_cols is not None: + features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + features_mask[adj_sample_skip_cols] = False + features_batch = features_batch[features_mask] + + torch.cuda.nvtx.range_pop() # nvtx-selectfeats + if epoch >= 1: + model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + # if epoch >= 1: + # start_time(start_inner_timer) + # torch.cuda.nvtx.range_push("nvtx-fakemats") + # features_batch = torch.cuda.FloatTensor(142247, 128) # SAGEConvfake + # adjs0_rows = torch.randint(0, 49387, (168660,)).cuda().long() # SAGEConvfake + # adjs0_cols = torch.randint(0, 142247, (168660,)).cuda().long() # SAGEConvfake + # adjs[0] = torch.stack((adjs0_rows, adjs0_cols)) # SAGEConvfake + # adjs1_rows = torch.randint(0, 8953, (57667,)).cuda().long() # SAGEConvfake + # adjs1_cols = torch.randint(0, 49387, (57667,)).cuda().long() # SAGEConvfake + # adjs[1] = torch.stack((adjs1_rows, adjs1_cols)) # SAGEConvfake + # adjs2_rows = torch.randint(0, 1024, (8125,)).cuda().long() # SAGEConvfake + # adjs2_cols = torch.randint(0, 8953, (8125,)).cuda().long() # SAGEConvfake + # adjs[2] = torch.stack((adjs2_rows, adjs2_cols)) # SAGEConvfake + + # features_batch = torch.cuda.FloatTensor(69740, 128) # GCNConvfake + # adjs0_rows = torch.randint(0, 9862, (91518,)).cuda().long() # GCNConvfake + # adjs0_cols = torch.randint(0, 69740, (91518,)).cuda().long() # GCNConvfake + # adjs0_indices = torch.stack((adjs0_rows, adjs0_cols)) # GCNConvfake + # adjs0_vals = torch.cuda.FloatTensor(91518).fill_(1.0) # GCNConvfake + # adjs[0] = torch.sparse_coo_tensor(adjs0_indices, adjs0_vals, torch.Size([9862, 69740])) # GCNConvfake + # adjs[0] = adjs[0].coalesce() + # adjs1_rows = torch.randint(0, 1024, (9065,)).cuda().long() # GCNConvfake + # adjs1_cols = torch.randint(0, 9862, (9065,)).cuda().long() # GCNConvfake + # adjs1_indices = torch.stack((adjs1_rows, adjs1_cols)) # GCNConvfake + # adjs1_vals = torch.cuda.FloatTensor(9065).fill_(1.0) # GCNConvfake + # adjs[1] = torch.sparse_coo_tensor(adjs1_indices, adjs1_vals, torch.Size([1024, 9862])) # GCNConvfake + # adjs[1] = adjs[1].coalesce() + # if epoch >= 1: + # model.timings["fakemats"].append(stop_time(start_inner_timer, stop_inner_timer)) + # torch.cuda.nvtx.range_pop() # nvtx-fakemats + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("nvtx-fwd") + # feat_dims.append(features_batch.size(0)) + # adj0_dims.append(adjs[0].size(1)) + # adj1_dims.append(adjs[1].size(1)) + # adj0_nnzs.append(adjs[0]._nnz()) + # adj1_nnzs.append(adjs[1]._nnz()) + + # seeds = next(train_loader_iter) + # n_id, batch_size, adjs = quiver_sampler.sample(batch_vtxs) + # adjs = [adj.to(rank) for adj in adjs] + # features_batch = features_loc[n_id] + optimizer.zero_grad() + logits = model(adjs, features_batch, epoch) + torch.cuda.nvtx.range_pop() # nvtx-fwd + if epoch >= 1: + model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-loss") + # loss = F.nll_loss(logits, data.y[batch_vtxs].long()) # GCNConv + loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs].long()) # SAGEConv + # loss = F.nll_loss(logits[:args.batch_size], data.y[seeds].long()) # quiver + + torch.cuda.nvtx.range_pop() # nvtx-loss + if epoch >= 1: + model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("nvtx-bwd") + loss.backward() + torch.cuda.nvtx.range_pop() # nvtx-bwd + if epoch >= 1: + model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + for W in model.parameters(): + dist.all_reduce(W.grad) + W.grad /= size + torch.cuda.nvtx.range_push("nvtx-optstep") + optimizer.step() + + for i in range(len(adjs)): + del adjs[0] + del adjs + torch.cuda.nvtx.range_pop() # nvtx-optstep + if epoch >= 1: + model.timings["train"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + + if last_batch: + # break + last_batch = False + args.n_bulkmb = tmp_bulkmb + args.batch_size = tmp_batch_size + + torch.cuda.nvtx.range_pop() # nvtx-training + for i in range(len(frontiers_bulk)): + del frontiers_bulk[0] + del frontiers_bulk + for i in range(len(adj_matrices_bulk)): + del adj_matrices_bulk[0] + del adj_matrices_bulk + torch.cuda.nvtx.range_pop() # nvtx-bulk + + dist.barrier() + if epoch >= 1: + dur.append(time.time() - epoch_start) + if rank == 0 and epoch >= 1 and (epoch % 5 == 0 or epoch == args.n_epochs - 1): + print("Evaluating", flush=True) + # acc = model.evaluate(g_loc, features_loc, labels_rank, rank_val_nids, \ + # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) + acc1 = 0.0 + acc3 = 0.0 + if True or args.timing: + sample_dur = [x / 1000 for x in model.timings["sample"]] + train_dur = [x / 1000 for x in model.timings["train"]] + fwd_dur = [x / 1000 for x in model.timings["fwd"]] + bwd_dur = [x / 1000 for x in model.timings["bwd"]] + selectfeats_dur = [x / 1000 for x in model.timings["selectfeats"]] + precomp_dur = [x / 1000 for x in model.timings["precomp"]] + spmm_dur = [x / 1000 for x in model.timings["spmm"]] + gemmi_dur = [x / 1000 for x in model.timings["gemm_i"]] + gemmw_dur = [x / 1000 for x in model.timings["gemm_w"]] + aggr_dur = [x / 1000 for x in model.timings["aggr"]] + fakemats_dur = [x / 1000 for x in model.timings["fakemats"]] + # sf_precomp_dur = [x / 1000 for x in model.timings["selectfeats-precomp"]] + # sf_sortdst_dur = [x / 1000 for x in model.timings["selectfeats-sortdst"]] + # sf_tallysplit_dur = [x / 1000 for x in model.timings["selectfeats-tallysplit"]] + # sf_tallya2a_dur = [x / 1000 for x in model.timings["selectfeats-tallya2a"]] + # sf_vtxsa2a_dur = [x / 1000 for x in model.timings["selectfeats-vtxsa2a"]] + # sf_featsa2a_dur = [x / 1000 for x in model.timings["selectfeats-featsa2a"]] + + # extract_dur = [x / 1000 for x in model.timings["extract"]] + extract_select_dur = [x / 1000 for x in model.timings["extract-select"]] + extract_inst_dur = [x / 1000 for x in model.timings["extract-inst"]] + extract_coalesce_dur = [x / 1000 for x in model.timings["extract-coalesce"]] + + print(f"sample: {np.sum(sample_dur) / epoch} extract-select: {np.sum(extract_select_dur) / epoch} extract-inst: {np.sum(extract_inst_dur) / epoch} extract-coalesce: {np.sum(extract_coalesce_dur) / epoch} train: {np.sum(train_dur) / epoch}", flush=True) + print(f"feats: {np.sum(selectfeats_dur) / epoch} fwd: {np.sum(fwd_dur) / epoch} bwd: {np.sum(bwd_dur) / epoch}") + print(f"precomp: {np.sum(precomp_dur)} spmm: {np.sum(spmm_dur)} gemmi: {np.sum(gemmi_dur)} gemmw: {np.sum(gemmw_dur)} aggr: {np.sum(aggr_dur)} fakemats: {np.sum(fakemats_dur)}") + # print(f"feat_dim.median: {np.median(feat_dims)}") + # print(f"adj0_dim.median: {np.median(adj0_dims)}") + # print(f"adj1_dim.median: {np.median(adj1_dims)}") + # print(f"adj0_nnz.median: {np.median(adj0_nnzs)}") + # print(f"adj1_nnz.median: {np.median(adj1_nnzs)}") + # print(f"adj_nnzs.median: {np.median(adj_nnzs)}") + print(f"total: {np.sum(dur) / epoch}", flush=True) + if False and args.sample_method == "sage" and args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): + model = model.cpu() + train_nid = train_nid.cpu() + test_nid = test_nid.cpu() + out = model.evaluate(adj_matrix, inputs, test_nid, data.y) + # correct_count = model.evaluate(adj_matrix, inputs, test_nid, data.y.cpu()) + res = out.argmax(dim=-1) == data.y.cpu() + # acc1 = int(res[train_nid].sum()) / train_nid.size(0) + acc3 = int(res[test_nid].sum()) / test_nid.size(0) + # acc3 = correct_count / test_nid.size(0) + train_nid = train_nid.to(device) + test_nid = test_nid.to(device) + model = model.to(device) + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) + else: + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), acc3), flush=True) + + dist.barrier() + + total_stop = time.time() + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='GCN') + parser.add_argument("--dataset", type=str, default="Cora", + help="dataset to train") + parser.add_argument("--sample-method", type=str, default="ladies", + help="sampling algorithm for training") + parser.add_argument("--dropout", type=float, default=0.5, + help="dropout probability") + parser.add_argument("--gpu", type=int, default=4, + help="gpus per node") + parser.add_argument("--lr", type=float, default=1e-2, + help="learning rate") + parser.add_argument("--n-epochs", type=int, default=200, + help="number of training epochs") + parser.add_argument("--n-hidden", type=int, default=256, + help="number of hidden gcn units") + parser.add_argument("--n-layers", type=int, default=1, + help="number of hidden gcn layers") + parser.add_argument("--batch-size", type=int, default=512, + help="number of vertices in minibatch") + parser.add_argument("--samp-num", type=str, default="2-2", + help="number of vertices per layer of minibatch") + parser.add_argument("--weight-decay", type=float, default=5e-4, + help="Weight for L2 loss") + parser.add_argument("--aggr", type=str, default="mean", + help="Aggregator type: mean/sum") + parser.add_argument('--world-size', default=-1, type=int, + help='number of nodes for distributed training') + parser.add_argument('--rank', default=-1, type=int, + help='node rank for distributed training') + parser.add_argument('--dist-url', default='env://', type=str, + help='url used to set up distributed training') + parser.add_argument('--dist-backend', default='nccl', type=str, + help='distributed backend') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument('--normalize', action="store_true", + help='normalize adjacency matrix') + parser.add_argument('--partitioning', default='ONE5D', type=str, + help='partitioning strategy to use') + parser.add_argument('--replication', default=1, type=int, + help='partitioning strategy to use') + parser.add_argument('--n-bulkmb', default=1, type=int, + help='number of minibatches to sample in bulk') + parser.add_argument('--bulk-batch-fetch', default=1, type=int, + help='number of minibatches to fetch features for in bulk') + parser.add_argument('--n-darts', default=-1, type=int, + help='number of darts to throw per minibatch in LADIES sampling') + parser.add_argument('--semibulk', default=128, type=int, + help='number of batches to column extract from in bulk') + parser.add_argument('--timing', action="store_true", + help='whether to turn on timers') + parser.add_argument('--baseline', action="store_true", + help='whether to avoid col selection for baseline comparison') + parser.add_argument('--replicate-graph', action="store_true", + help='replicate adjacency matrix on each device') + args = parser.parse_args() + args.samp_num = [int(i) for i in args.samp_num.split('-')] + print(args) + + main(args) diff --git a/examples/gcn_1d.py b/examples/gcn_1d.py new file mode 100644 index 0000000..f12919c --- /dev/null +++ b/examples/gcn_1d.py @@ -0,0 +1,380 @@ +import argparse +import math +import os +import os.path as osp +import time +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.autograd.profiler as profiler +import torch_geometric +from torch_geometric.data import Data +from torch_geometric.datasets import Planetoid, Reddit +import torch_geometric.transforms as T +from torch_geometric.utils import add_remaining_self_loops +import torch_sparse + +from cagnet.nn.conv import GCNConv +from cagnet.partitionings import Partitioning +import cagnet.nn.functional as CAGF + +import socket + +class GCN(nn.Module): + def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, partitioning, device, + group=None, row_groups=None, col_groups=None): + super(GCN, self).__init__() + self.layers = nn.ModuleList() + self.rank = rank + self.size = size + self.group = group + self.device = device + self.partitioning = partitioning + self.timings = dict() + + # input layer + self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # hidden layers + for i in range(n_layers - 1): + self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # output layer + self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) + + def forward(self, graph, inputs, ampbyp): + h = inputs + for l, layer in enumerate(self.layers): + h = layer(self, graph, h, ampbyp) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + + h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + return h + +# Normalize all elements according to KW's normalization rule +def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx, normalization): + if not normalization: + return adj_part + + adj_part = adj_part.coalesce() + deg = torch.histc(adj_matrix[0].double(), bins=node_count) + deg = deg.pow(-0.5) + + row_len = adj_part.size(0) + col_len = adj_part.size(1) + + dleft = torch.sparse_coo_tensor([np.arange(0, row_len).tolist(), + np.arange(0, row_len).tolist()], + deg[row_vtx:(row_vtx + row_len)].float(), + size=(row_len, row_len), + requires_grad=False, device=torch.device("cpu")) + + dright = torch.sparse_coo_tensor([np.arange(0, col_len).tolist(), + np.arange(0, col_len).tolist()], + deg[col_vtx:(col_vtx + col_len)].float(), + size=(col_len, col_len), + requires_grad=False, device=torch.device("cpu")) + # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) + ad_ind, ad_val = torch_sparse.spspmm(adj_part._indices(), adj_part._values(), + dright._indices(), dright._values(), + adj_part.size(0), adj_part.size(1), dright.size(1)) + + adj_part_ind, adj_part_val = torch_sparse.spspmm(dleft._indices(), dleft._values(), + ad_ind, ad_val, + dleft.size(0), dleft.size(1), adj_part.size(1)) + + adj_part = torch.sparse_coo_tensor(adj_part_ind, adj_part_val, + size=(adj_part.size(0), adj_part.size(1)), + requires_grad=False, device=torch.device("cpu")) + + return adj_part + +# Split a COO into partitions of size n_per_proc +# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. +def split_coo(adj_matrix, node_count, n_per_proc, dim): + vtx_indices = list(range(0, node_count, n_per_proc)) + vtx_indices.append(node_count) + + am_partitions = [] + for i in range(len(vtx_indices) - 1): + am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] + am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] + am_part[dim] -= vtx_indices[i] + am_partitions.append(am_part) + + return am_partitions, vtx_indices + +def oned_partition(rank, size, inputs, adj_matrix, data, features, classes, device, normalize): + node_count = inputs.size(0) + n_per_proc = math.ceil(float(node_count) / size) + + am_partitions = None + am_pbyp = None + + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) + + # Compute the adj_matrix and inputs partitions for this process + # TODO: Maybe I do want grad here. Unsure. + with torch.no_grad(): + # Column partitions + am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) + + proc_node_count = vtx_indices[rank + 1] - vtx_indices[rank] + am_pbyp, _ = split_coo(am_partitions[rank], node_count, n_per_proc, 0) + for i in range(len(am_pbyp)): + if i == size - 1: + last_node_count = vtx_indices[i + 1] - vtx_indices[i] + am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), + size=(last_node_count, proc_node_count), + requires_grad=False) + + am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + vtx_indices[rank], normalize) + else: + am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), + size=(n_per_proc, proc_node_count), + requires_grad=False) + + am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + vtx_indices[rank], normalize) + + for i in range(len(am_partitions)): + proc_node_count = vtx_indices[i + 1] - vtx_indices[i] + am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], + torch.ones(am_partitions[i].size(1)), + size=(node_count, proc_node_count), + requires_grad=False) + am_partitions[i] = scale_elements(adj_matrix, am_partitions[i], node_count, 0, vtx_indices[i], normalize) + + input_partitions = torch.split(inputs, math.ceil(float(inputs.size(0)) / size), dim=0) + + adj_matrix_loc = am_partitions[rank] + inputs_loc = input_partitions[rank] + + print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) + print(f"rank: {rank} inputs.size: {inputs.size()}", flush=True) + return inputs_loc, adj_matrix_loc, am_pbyp + +def evaluate(model, graph, features, labels, nid, nid_count, ampbyp, group): + model.eval() + with torch.no_grad(): + non_eval_timings = copy.deepcopy(model.timings) + logits = model(graph, features, ampbyp, degrees) + model.timings = non_eval_timings # don't include evaluation timings + + # # all-gather logits across ranks + # logits_recv = [] + # for i in range(len(ampbyp)): + # logits_recv.append(torch.FloatTensor(ampbyp[0].size(1), logits.size(1))) + + # if logits.size(0) != ampbyp[0].size(1): + # pad_row = ampbyp[0].size(1) - logits.size(0) + # logits = torch.cat((logits, torch.FloatTensor(pad_row, logits.size(1))), dim=0) + + # dist.all_gather(logits_recv, logits, group) + + # padding = graph.size(0) - ampbyp[0].size(1) * (len(ampbyp) - 1) + # logits_recv[-1] = logits_recv[-1][:padding,:] + + # logits = torch.cat(logits_recv) + + logits = logits[nid] + labels = labels[nid] + if logits.size(0) > 0: + _, indices = torch.max(logits, dim=1) + correct = torch.sum(indices == labels) + else: + correct = torch.tensor(0) + + dist.all_reduce(correct, op=dist.reduce_op.SUM, group=group) + return correct.item() * 1.0 / nid_count + # return correct.item() * 1.0 / len(labels) + +def main(args): + # load and preprocess dataset + device = torch.device('cuda') + + path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', args.dataset) + + if args.dataset == "Cora" or args.dataset == "Reddit": + if args.dataset == "Cora": + dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures()) + elif args.dataset == "Reddit": + dataset = Reddit(path, transform=T.NormalizeFeatures()) + + data = dataset[0] + data = data.to(device) + data.x.requires_grad = True + inputs = data.x.to(device) + inputs.requires_grad = True + data.y = data.y.to(device) + edge_index = data.edge_index + num_features = dataset.num_features + num_classes = dataset.num_classes + adj_matrix = edge_index + elif args.dataset == "Amazon": + print(f"Loading coo...", flush=True) + edge_index = torch.load("../../data/Amazon/processed/data.pt") + print(f"Done loading coo", flush=True) + n = 14249639 + num_features = 300 + num_classes = 24 + inputs = torch.rand(n, num_features) + data = Data() + data.y = torch.rand(n).uniform_(0, num_classes - 1).long() + data.train_mask = torch.ones(n).long() + adj_matrix = edge_index.t_() + data = data.to(device) + inputs.requires_grad = True + data.y = data.y.to(device) + + # Initialize distributed environment with SLURM + if "SLURM_PROCID" in os.environ.keys(): + os.environ["RANK"] = os.environ["SLURM_PROCID"] + + if "SLURM_NTASKS" in os.environ.keys(): + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + + os.environ["MASTER_ADDR"] = args.hostname + os.environ["MASTER_PORT"] = "1234" + + dist.init_process_group(backend=args.dist_backend) + rank = dist.get_rank() + size = dist.get_world_size() + print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) + + if args.normalize: + adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + + partitioning = Partitioning.ONED + + print(f"partitioning...", flush=True) + group = dist.new_group(list(range(size))) + features_loc, g_loc, ampbyp = oned_partition(rank, size, inputs, adj_matrix, data, \ + inputs, num_classes, device, args.normalize) + print(f"done partitioning", flush=True) + + features_loc = features_loc.to(device) + g_loc = g_loc.to(device) + for i in range(len(ampbyp)): + ampbyp[i] = ampbyp[i].t().coalesce().to(device) + + + # create GCN model + torch.manual_seed(0) + model = GCN(num_features, + args.n_hidden, + num_classes, + args.n_layers, + rank, + size, + partitioning, + device, + group) + + # use optimizer + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + n_per_proc = math.ceil(float(inputs.size(0)) / size) + + rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank] + labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank] + rank_train_nids = rank_train_mask.nonzero().squeeze() + + train_nid = data.train_mask.nonzero().squeeze() + + torch.manual_seed(0) + total_start = time.time() + for epoch in range(args.n_epochs): + print(f"Epoch: {epoch}", flush=True) + if epoch == 1: + total_start = time.time() + model.train() + + # forward + logits = model(g_loc, features_loc, ampbyp) + loss = CAGF.cross_entropy(logits[rank_train_nids], labels_rank[rank_train_nids], train_nid.size(0), \ + num_classes, partitioning, rank, group, size) + + optimizer.zero_grad() + loss.backward() + + optimizer.step() + + """ + # acc = evaluate(model, g_loc, features_loc, labels, val_nid, ampbyp, ampbyp_dgl, degrees, col_groups[0]) + # acc = evaluate(model, g_loc, features_loc, labels, val_nid, \ + acc = evaluate(model, g_loc, features_loc, labels_rank, rank_val_nids, \ + val_mask.nonzero().squeeze().size(0), ampbyp, ampbyp_dgl, degrees, col_groups[0]) + print("Rank: {:05d} | Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | " + "ETputs(KTEPS) {:.2f}".format(rank, epoch, np.mean(dur), loss.item(), + acc, n_edges / np.mean(dur) / 1000), flush=True) + """ + dist.barrier() + total_stop = time.time() + print(f"rank: {rank} total_time: {total_stop - total_start}") + + """ + # print(prof.key_averages().table(sort_by='self_cpu_time_total', row_limit=10)) + barrier_start = time.time() + dist.barrier() + model.timings["total"] += time.time() - total_start + model.timings["barrier"] += time.time() - barrier_start + print() + # # acc = evaluate(model, g_loc, features_loc, labels, test_nid, ampbyp, ampbyp_dgl, degrees, col_groups[0]) + # acc = evaluate(model, g_loc, features_loc, labels, test_nid, \ + acc = evaluate(model, g_loc, features_loc, labels_rank, rank_test_nids, \ + test_mask.nonzero().squeeze().size(0), ampbyp, ampbyp_dgl, degrees, col_groups[0]) + print("Test Accuracy {:.4f}".format(acc)) + + print(flush=True) + dist.barrier() + print("Timings") + model.timings["comp"] = model.timings["scomp"] + model.timings["dcomp"] + model.timings["comm"] = model.timings["bcast"] + model.timings["reduce"] + model.timings["op"] + # print("rank, total, scomp, dcomp, bcast, reduce, op, barrier") + # print(f"{rank}, {model.timings['total']}, {model.timings['scomp']}, {model.timings['dcomp']}, {model.timings['bcast']}, {model.timings['reduce']}, {model.timings['op']}, {model.timings['barrier']}") + print(f"rank: {rank} timings: {model.timings}") + """ + print(f"rank: {rank} {logits}") + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='GCN') + parser.add_argument("--dataset", type=str, default="Cora", + help="dataset to train") + parser.add_argument("--dropout", type=float, default=0.5, + help="dropout probability") + parser.add_argument("--gpu", type=int, default=-1, + help="gpu") + parser.add_argument("--lr", type=float, default=1e-2, + help="learning rate") + parser.add_argument("--n-epochs", type=int, default=200, + help="number of training epochs") + parser.add_argument("--n-hidden", type=int, default=16, + help="number of hidden gcn units") + parser.add_argument("--n-layers", type=int, default=1, + help="number of hidden gcn layers") + parser.add_argument("--weight-decay", type=float, default=5e-4, + help="Weight for L2 loss") + parser.add_argument("--aggregator-type", type=str, default="gcn", + help="Aggregator type: mean/gcn/pool/lstm") + parser.add_argument('--world-size', default=-1, type=int, + help='number of nodes for distributed training') + parser.add_argument('--rank', default=-1, type=int, + help='node rank for distributed training') + parser.add_argument('--dist-url', default='env://', type=str, + help='url used to set up distributed training') + parser.add_argument('--dist-backend', default='nccl', type=str, + help='distributed backend') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument('--normalize', action="store_true", + help='normalize adjacency matrix') + parser.add_argument('--partitioning', default='ONED', type=str, + help='partitioning strategy to use') + args = parser.parse_args() + print(args) + + main(args) diff --git a/examples/gnn_train_ctd.yaml b/examples/gnn_train_ctd.yaml new file mode 100644 index 0000000..95ccc3c --- /dev/null +++ b/examples/gnn_train_ctd.yaml @@ -0,0 +1,77 @@ +stage: edge_classifier +model: InteractionGNN2 +input_dir: /global/cfs/cdirs/m4439/CTD2023_results/module_map/ # Should contain the files used in training and inference +stage_dir: /global/cfs/cdirs/m1982/alokt/data/trackml/ctd/gnn # A directory to be created that will store logs, artifacts, and output data +project: CTD_2023 # Used in logging +accelerator: gpu +devices: 1 +nodes: 1 + +# Dataset parameters +data_split: [7800, 1000, 1000] # Number of [training, validation, testing] examples +dataset_class: GraphDataset +undirected: False + +# Truth and weighting parameters. Syntax is... +# Single value: track[key] == value +# List of 2 floats: value1 <= track[key] <= value2 +# List with a string as the first value: First value defines the boolean test, the second value is a list or value. Boolean tests are: (in, not_in, is, is_not, within, not_within) +# All other tracks follow the default: If a true edge, weight=0. If a fake edge, weight=1. +weighting: + - weight: 0.1 + conditions: + y: False + - weight: 0. + conditions: + y: True + - weight: 1. + conditions: + y: True + pt: [1000, .inf] + nhits: [3, .inf] + primary: True + pdgId: [not_in, [11, -11]] + radius: [0., 260.] + eta_particle: [-4., 4.] + redundant_split_edges: False + +# Model parameters +edge_cut: 0.5 +node_features: [r, phi, z, eta, cluster_r_1, cluster_phi_1, cluster_z_1, cluster_eta_1, cluster_r_2, cluster_phi_2, cluster_z_2, cluster_eta_2] +node_scales: [1000., 3.14159265359, 1000., 1., 1000., 3.14159265359, 1000., 1., 1000., 3.14159265359, 1000., 1.] +edge_features: [dr, dphi, dz, deta, phislope, rphislope] +#edge_features: None +hidden: 128 +n_graph_iters: 8 +n_node_encoder_layers: 3 +n_edge_encoder_layers: 3 +n_node_net_layers: 3 +n_edge_net_layers: 3 +n_node_decoder_layers: 3 +n_edge_decoder_layers: 3 +layernorm: False +output_layer_norm: False +edge_output_transform_final_layer_norm: False +batchnorm: True +output_batch_norm: False +edge_output_transform_final_batch_norm: False +bn_track_running_stats: False +hidden_activation: ReLU +output_activation: ReLU +edge_output_transform_final_activation: null +concat: True +node_net_recurrent: False +edge_net_recurrent: False +in_out_diff_agg: True +checkpointing: True + +# Training parameters +warmup: 5 +lr: 0.0005 +min_lr: 0.000005 +factor: 0.9 +patience: 15 +max_epochs: 1000 +debug: false +num_workers: [8,8,8] +max_training_graph_size: 2800000 diff --git a/examples/gnn_train_physics_ex3.yaml b/examples/gnn_train_physics_ex3.yaml new file mode 120000 index 0000000..ecfb8c8 --- /dev/null +++ b/examples/gnn_train_physics_ex3.yaml @@ -0,0 +1 @@ +../../acorn/examples/Example_3/gnn_train.yaml \ No newline at end of file diff --git a/examples/interaction_gnn_15d.py b/examples/interaction_gnn_15d.py new file mode 100644 index 0000000..bc83eb6 --- /dev/null +++ b/examples/interaction_gnn_15d.py @@ -0,0 +1,1615 @@ +import argparse +import copy +import math +import os +import os.path as osp +import time +import torch.multiprocessing as mp +import numpy as np + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.autograd.profiler as profiler +import torch_geometric +import torch_geometric.transforms as T +import torch_sparse + +from torch_geometric.data import Batch, Data, Dataset +from torch_geometric.nn import SAGEConv +from torch_geometric.datasets import Planetoid, Reddit +from torch_geometric.loader import DataLoader, NeighborSampler, NeighborLoader +from torch_geometric.loader.shadow import ShaDowKHopSampler +# from shadow import ShaDowKHopSampler +from torch_geometric.transforms import LargestConnectedComponents +from torch_geometric.typing import SparseTensor +from torch_geometric.utils import add_remaining_self_loops +from torch_scatter import scatter_add +from torch.multiprocessing import Process +from torchmetrics.classification import Precision, Recall, F1Score, ROC, AUROC, Accuracy +from torch.nn.parallel import DistributedDataParallel +from torch.utils.data.distributed import DistributedSampler +from torch.utils.checkpoint import checkpoint + +from cagnet.nn.conv import GCNConv +from cagnet.partitionings import Partitioning +from cagnet.samplers import shadow_sampler +from cagnet.samplers.utils import * +import cagnet.nn.functional as CAGF +import torch.nn.functional as F + +from datetime import timedelta +from sparse_coo_tensor_cpp import sort_dst_proc_gpu +from sklearn.metrics import roc_auc_score +from statistics import median + +import socket +import wandb +import yaml + +class InteractionGNN(nn.Module): + def __init__(self, in_feats, n_hidden, n_classes, + nb_node_layer, nb_edge_layer, n_graph_iters, + node_features, edge_features, layernorm, batchnorm, + hidden_activation, output_activation, + aggr, checkpointing, replication, device, + impl="cagnet", dataset="physics_ex3"): + + super(InteractionGNN, self).__init__() + # self.layers = nn.ModuleList() + self.nb_node_layer = nb_node_layer + self.nb_edge_layer = nb_edge_layer + self.n_graph_iters = n_graph_iters + self.node_features = node_features + self.edge_features = edge_features + self.layernorm = layernorm + self.batchnorm = batchnorm + self.hidden_activation = hidden_activation + self.output_activation = output_activation + self.aggr = aggr + self.checkpointing = checkpointing + self.replication = replication + self.timings = dict() + self.impl = impl + self.dataset = dataset + self.device = device + + self.timings["sample"] = [] + self.timings["train"] = [] + self.timings["extract"] = [] + self.timings["fwdbwd"] = [] + self.timings["fwd"] = [] + self.timings["bwd"] = [] + self.timings["h2d"] = [] + self.timings["allreduce"] = [] + self.timings["allreduce-pt1"] = [] + self.timings["allreduce-pt2"] = [] + self.timings["allreduce-pt3"] = [] + self.timings["allreduce-pt4"] = [] + self.timings["shadow-sampling"] = [] + self.timings["shadow-selection"] = [] + self.timings["shadow-extraction"] = [] + self.timings["fwd-prelude"] = [] + self.timings["fwd-encoder"] = [] + self.timings["fwd-graphiters"] = [] + + torch.manual_seed(1) + # aggr_list = ["sum", "mean", "max", "std"] + # self.aggregation = torch_geometric.nn.aggr.MultiAggregation(aggr_list, mode="cat") + + # network_input_size = (1 + 2 * len(aggr_list)) * n_hidden + + self.node_encoder = self.make_mlp( + input_size=in_feats, + sizes=[n_hidden] * nb_node_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + + if len(self.edge_features) > 0: + self.edge_encoder = self.make_mlp( + input_size=len(self.edge_features), + sizes=[n_hidden] * nb_edge_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + else: + self.edge_encoder = self.make_mlp( + input_size=2 * n_hidden, + sizes=[n_hidden] * nb_edge_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + + torch.manual_seed(1) + # aggr_list = ["sum", "mean", "max", "std"] + # self.aggregation = torch_geometric.nn.aggr.MultiAggregation(aggr_list, mode="cat") + + # network_input_size = (1 + 2 * len(aggr_list)) * n_hidden + + self.node_encoder = self.make_mlp( + input_size=in_feats, + sizes=[n_hidden] * nb_node_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + + if len(self.edge_features) > 0: + self.edge_encoder = self.make_mlp( + input_size=len(self.edge_features), + sizes=[n_hidden] * nb_edge_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + else: + self.edge_encoder = self.make_mlp( + input_size=2 * n_hidden, + sizes=[n_hidden] * nb_edge_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + + in_edge_net = n_hidden * 6 + self.edge_network = nn.ModuleList( + [ + self.make_mlp( + input_size=in_edge_net, + sizes=[n_hidden] * nb_edge_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + for i in range(n_graph_iters) + ] + ) + + in_node_net = n_hidden * 4 + self.node_network = nn.ModuleList( + [ + self.make_mlp( + input_size=in_node_net, + sizes=[n_hidden] * nb_node_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + for i in range(n_graph_iters) + ] + ) + + # edge decoder + self.edge_decoder = self.make_mlp( + input_size=n_hidden, + sizes=[n_hidden] * nb_edge_layer, + output_activation=output_activation, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + + # edge output transform layer + self.edge_output_transform = self.make_mlp( + input_size=n_hidden, + sizes=[n_hidden, 1], + output_activation=None, + hidden_activation=hidden_activation, + layer_norm=layernorm, + batch_norm=batchnorm, + ) + + # dropout layer + self.dropout = nn.Dropout(p=0.1) + + def make_mlp( + self, + input_size, + sizes, + hidden_activation="ReLU", + output_activation=None, + layer_norm=False, # TODO : change name to hidden_layer_norm while ensuring backward compatibility + output_layer_norm=False, + batch_norm=False, # TODO : change name to hidden_batch_norm while ensuring backward compatibility + output_batch_norm=False, + input_dropout=0, + hidden_dropout=0, + track_running_stats=False, + ): + """Construct an MLP with specified fully-connected layers.""" + hidden_activation = getattr(nn, hidden_activation) + if output_activation is not None: + output_activation = getattr(nn, output_activation) + layers = [] + n_layers = len(sizes) + sizes = [input_size] + sizes + # Hidden layers + for i in range(n_layers - 1): + if i == 0 and input_dropout > 0: + layers.append(nn.Dropout(input_dropout)) + layers.append(nn.Linear(sizes[i], sizes[i + 1])) + if layer_norm: # hidden_layer_norm + layers.append(nn.LayerNorm(sizes[i + 1], elementwise_affine=False)) + if batch_norm: # hidden_batch_norm + layers.append( + nn.BatchNorm1d( + sizes[i + 1], + eps=6e-05, + track_running_stats=track_running_stats, + affine=True, + ) # TODO : Set BatchNorm and LayerNorm parameters in config file ? + ) + layers.append(hidden_activation()) + if hidden_dropout > 0: + layers.append(nn.Dropout(hidden_dropout)) + # Final layer + layers.append(nn.Linear(sizes[-2], sizes[-1])) + if output_activation is not None: + if output_layer_norm: + layers.append(nn.LayerNorm(sizes[-1], elementwise_affine=False)) + if output_batch_norm: + layers.append( + nn.BatchNorm1d( + sizes[-1], + eps=6e-05, + track_running_stats=track_running_stats, + affine=True, + ) # TODO : Set BatchNorm and LayerNorm parameters in config file ? + ) + layers.append(output_activation()) + return nn.Sequential(*layers) + + def message_step(self, x, e, src, dst, i): + # Compute new node features + edge_inputs = torch.cat([e, x[src], x[dst]], dim=-1) # order dst src x ? + e_updated = self.edge_network[i](edge_inputs) + edge_messages_from_src = scatter_add(e_updated, dst, dim=0, dim_size=x.shape[0]) + edge_messages_from_dst = scatter_add(e_updated, src, dim=0, dim_size=x.shape[0]) + + node_inputs = torch.cat( + [edge_messages_from_src, edge_messages_from_dst, x], dim=-1 + ) # to check : the order dst src x ? + x_updated = self.node_network[i](node_inputs) + + return ( + x_updated, + e_updated, + self.edge_output_transform(self.edge_decoder(e_updated)), + ) + + def output_step(self, x, start, end, e): + classifier_inputs = torch.cat([x[start], x[end], e], dim=1) + classifier_output = self.output_edge_classifier(classifier_inputs).squeeze(-1) + return classifier_output + + def forward(self, batch, epoch=0): + # x = torch.stack([batch["z"]], dim=-1).float() + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + if epoch >= 1: + start_time(start_timer, timing_arg=True) + x = torch.stack( + [batch[feature] for feature in self.node_features], dim=-1 + ).float() + + if len(self.edge_features) > 0: + edge_attr = torch.stack( + [batch[feature] for feature in self.edge_features], dim=-1 + ).float() + else: + edge_attr = None + + if batch.edge_index is not None: + src, dst = batch.edge_index + elif batch.adj_t is not None: + coo = batch.adj_t.coo() + src, dst = coo[0], coo[1] + + x.requires_grad = True + if edge_attr is not None: + edge_attr.requires_grad = True + + if epoch >= 1: + self.timings["fwd-prelude"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + + if epoch >= 1: + start_time(start_timer, timing_arg=True) + if self.checkpointing: + x = checkpoint(self.node_encoder, x, use_reentrant=False) + if edge_attr is not None: + e = checkpoint(self.edge_encoder, edge_attr, use_reentrant=False) + else: + e = checkpoint(self.edge_encoder, torch.cat([x[src], x[dst]], dim=-1), use_reentrant=False) + else: + x = self.node_encoder(x) + if edge_attr is not None: + e = self.edge_encoder(edge_attr) + else: + e = self.edge_encoder(torch.cat([x[src], x[dst]], dim=-1)) + if epoch >= 1: + self.timings["fwd-encoder"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + + # if concat + if epoch >= 1: + start_time(start_timer, timing_arg=True) + input_x = x + input_e = e + # outputs = [] + outputs = None + for i in range(self.n_graph_iters): + # if concat + if self.checkpointing: + x = torch.cat([x, input_x], dim=-1) + e = torch.cat([e, input_e], dim=-1) + x, e, out = checkpoint(self.message_step, x, e, src, dst, i, use_reentrant=False) + # outputs.append(out) + outputs = out + else: + x = torch.cat([x, input_x], dim=-1) + e = torch.cat([e, input_e], dim=-1) + x, e, out = self.message_step(x, e, src, dst, i) + outputs = out + if epoch >= 1: + self.timings["fwd-graphiters"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + + # return outputs[-1].squeeze(-1) + return outputs.squeeze(-1) + + def loss_function(self, output, batch, balance="proportional"): + """ + Applies the loss function to the output of the model and the truth labels. + To balance the positive and negative contribution, simply take the means of each separately. + Any further fine tuning to the balance of true target, true background and fake can be handled + with the `weighting` config option. + """ + + assert hasattr(batch, "y"), "The batch does not have a truth label" + assert hasattr(batch, "weights"), "The batch does not have a weighting label" + + assert hasattr(batch, "y"), ( + "The batch does not have a truth label. Please ensure the batch has a `y`" + " attribute." + ) + assert hasattr(batch, "weights"), ( + "The batch does not have a weighting label. Please ensure the batch" + " weighting is handled in preprocessing." + ) + + if balance not in ["equal", "proportional"]: + warnings.warn( + f"{balance} is not a proper choice for the loss balance. Use either 'equal' or 'proportional'. Automatically switching to 'proportional' instead." + ) + balance = "proportional" + + negative_mask = ((batch.y == 0) & (batch.weights != 0)) | (batch.weights < 0) + + negative_loss = F.binary_cross_entropy_with_logits( + output[negative_mask], + torch.zeros_like(output[negative_mask]), + weight=batch.weights[negative_mask].abs(), + reduction="sum", + ) + + positive_mask = (batch.y == 1) & (batch.weights > 0) + positive_loss = F.binary_cross_entropy_with_logits( + output[positive_mask], + torch.ones_like(output[positive_mask]), + weight=batch.weights[positive_mask].abs(), + reduction="sum", + ) + + if balance == "proportional": + n = positive_mask.sum() + negative_mask.sum() + return ( + (positive_loss + negative_loss) / n, + positive_loss.detach() / n, + negative_loss.detach() / n, + ) + else: + n_pos, n_neg = positive_mask.sum(), negative_mask.sum() + n = n_pos + n_neg + return ( + positive_loss / n_pos + negative_loss / n_neg, + positive_loss.detach() / n, + negative_loss.detach() / n, + ) + + @torch.no_grad() + def evaluate(self, test_loader, epoch, curr_lr, epoch_time, wandb_arg, rank): + """ + The gateway for the evaluation stage. This class method is called from the eval_stage.py script. + """ + val_auc = AUROC(task="binary").to(self.device) + val_precision = Precision(task="binary").to(self.device) + val_recall = Recall(task="binary").to(self.device) + val_f1 = F1Score(task="binary").to(self.device) + val_loss = [] + eff = [] + tar_pur = [] + tot_pur = [] + + print("Evaluating") + test_batches = [] + for i, batch in enumerate(test_loader): + # batch = batch.to(self.device) + if rank == 0: + print(f"batch {i}") + # batch = batch.to(self.device) + batch = batch.cuda() + output = self(batch) + loss, pos_loss, neg_loss = self.loss_function(output, batch) + + scores = torch.sigmoid(output) + batch.scores = scores.detach() + test_batches.append(batch) + + all_truth = batch.y.bool() + target_truth = (batch.weights > 0) & all_truth + + + preds = torch.sigmoid(output) > 0.5 + + # Positives + edge_positive = preds.sum().float() + + # Signal true & signal tp + target_true = target_truth.sum().float() + target_true_positive = (target_truth.bool() & preds).sum().float() + all_true_positive = (all_truth.bool() & preds).sum().float() + target_auc = roc_auc_score( + target_truth.bool().cpu().detach(), torch.sigmoid(output).cpu().detach() + ) + # Eff, pur, auc + target_eff = target_true_positive / target_true + target_pur = target_true_positive / edge_positive + total_pur = all_true_positive / edge_positive + + val_losses = self.loss_function(output, batch) + val_losses = sum(val_losses) + val_loss.append(val_losses) + eff.append(target_eff.item()) + tar_pur.append(target_pur.item()) + tot_pur.append(total_pur.item()) + val_precision.update(preds, target_truth) + val_recall.update(preds, target_truth) + val_f1.update(preds, target_truth) + val_auc.update(output, target_truth) + + if rank == 0: + avg_loss = sum(val_loss) / len(val_loss) + target_eff = sum(eff) / len(eff) + avg_tarpur = sum(tar_pur) / len(tar_pur) + avg_totpur = sum(tot_pur) / len(tot_pur) + + efficiency = val_precision.compute() + purity = val_recall.compute() + f1_score = val_f1.compute() + auc = val_auc.compute() + print(f"Logging") + if rank == 0: + if wandb_arg: + wandb.log({ + # "train_loss": loss, + "current_lr": curr_lr, + "eff": efficiency, + # "target_pur": avg_tarpur, + "total_pur": purity, + "f1_score": f1_score, + "auc": auc, + "val_loss": avg_loss, + "epoch": epoch, + "time": epoch_time, + # "trainer/global_step": step + epoch*step + # "trainer/global_step": epoch + }) + # print(f"Target Purity: {avg_tarpur}", flush=True) + print(f"Total Purity: {purity}", flush=True) + print(f"Efficiency: {efficiency}", flush=True) + print(f"F1-Score: {f1_score}", flush=True) + print(f"AUC: {target_auc}", flush=True) + + # val_auc.reset(), val_precision.reset(), val_recall.reset() + + full_auc = 0.0 + masked_auc = 0.0 + # if self.dataset == "physics_ex3": + # filename = f"{self.impl}_{self.dataset}" + # full_auc, masked_auc = graph_roc_curve("testset", test_batches, "Interaction GNN ROC curve", filename) + # print(f"full_auc: {full_auc} type(full_auc): {type(full_auc)}", flush=True) + # print(f"masked_auc: {masked_auc} type(masked_auc): {type(masked_auc)}", flush=True) + # elif False and self.dataset == "ctd": + # filename = f"{self.impl}_{self.dataset}_edgewise" + # plot_config = {"title": "GNN edge-wise Efficiency vs (r,z)", + # "filename": filename } + # target_tracks = {"pt": [1000, np.inf], + # "nhits": [3, np.inf], + # "primary": True, + # # pdgId: {} + # "radius": [0., 260.], + # "eta_particle": [-4., 4.], + # "redundant_split_edges": False} + + # config = {"dataset": self.dataset, + # "score_cut": 0.5, + # "target_tracks": target_tracks, + # "stage_dir": filename } + # signal_efficiency = gnn_efficiency_rz(test_batches, plot_config, config) + # plot_config["vmin"] = 0.4 + # # gnn_purity_rz(test_batches, plot_config, config) + # gnn_purity_rz(test_loader, plot_config, config) + # full_auc = signal_efficiency + # masked_auc = 0.0 + # print(f"signal_efficiency: {signal_efficiency}", flush=True) + + return full_auc, masked_auc + + # # Load data from testset directory + # graph_constructor = cls(config).to(device) + # if checkpoint is not None: + # print(f"Restoring model from {checkpoint}") + # graph_constructor = cls.load_from_checkpoint(checkpoint, hparams=config).to( + # device + # ) + # graph_constructor.setup(stage="test") + + # plots: + # graph_scoring_efficiency: + # title: Interaction GNN Edge-wise Efficiency + # pt_units: GeV + # graph_roc_curve: + # title: Interaction GNN ROC curve + # all_plots = config["plots"] + + ## TODO: Handle the list of plots properly + #for plot_function, plot_config in all_plots.items(): + # if hasattr(eval_utils, plot_function): + # getattr(eval_utils, plot_function)( + # graph_constructor, plot_config, config + # ) + # else: + # print(f"Plot {plot_function} not implemented") + + def ignn_allreduce(self, params, size, epoch): + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + torch.cuda.synchronize() + if epoch >= 1: + start_time(start_timer, timing_arg=True) + grad_tensors = [] + tensor_names = [] + param_list = list(params) + for name, W in param_list: + if W.grad is not None: + # print(f"name: {name} W.grad.size(): {W.grad.size()} ndim: {W.grad.ndim}", flush=True) + if W.grad.size(0) == 64: + if W.grad.ndim == 2: + grad_tensors.append(W.grad) + tensor_names.append(name) + elif W.grad.ndim == 1: + grad_tensors.append(W.grad.unsqueeze(-1)) + tensor_names.append(name) + # else: + # print(f"tensor error: {name} W.grad.size: {W.grad.size()}") + + stacked_grad_tensors = torch.cat(grad_tensors, dim=1) + if epoch >= 1: + self.timings["allreduce-pt1"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + + stacked_grad_tensors = stacked_grad_tensors.contiguous() + torch.cuda.synchronize() + if epoch >= 1: + start_time(start_timer, timing_arg=True) + dist.all_reduce(stacked_grad_tensors) + if epoch >= 1: + self.timings["allreduce-pt2"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + + torch.cuda.synchronize() + if epoch >= 1: + start_time(start_timer, timing_arg=True) + stacked_grad_tensors /= size + grad_tensor_idx = 0 + name_tensor_idx = 0 + for name, W in param_list: + if W.grad is not None: + if W.grad.size(0) == 64: + name_tensor_idx += 1 + if W.grad.ndim == 2: + W.grad = stacked_grad_tensors[:,\ + grad_tensor_idx:grad_tensor_idx + W.grad.size(1)] + grad_tensor_idx += W.grad.size(1) + elif W.grad.ndim == 1: + W.grad = stacked_grad_tensors[:,\ + grad_tensor_idx:grad_tensor_idx + 1].squeeze() + grad_tensor_idx += 1 + if epoch >= 1: + self.timings["allreduce-pt3"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + + if epoch >= 1: + start_time(start_timer, timing_arg=True) + for _, W in params: + if W.grad is not None and W.grad.size(0) != 64: + dist.all_reduce(W.grad) + W.grad /= size + if epoch >= 1: + self.timings["allreduce-pt4"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + +def get_proc_groups(rank, size, replication): + rank_c = rank // replication + + row_procs = [] + for i in range(0, size, replication): + row_procs.append(list(range(i, i + replication))) + + col_procs = [] + for i in range(replication): + col_procs.append(list(range(i, size, replication))) + + row_groups = [] + for i in range(len(row_procs)): + row_groups.append(dist.new_group(row_procs[i])) + + col_groups = [] + for i in range(len(col_procs)): + col_groups.append(dist.new_group(col_procs[i])) + + return row_groups, col_groups + +# Normalize all elements according to KW's normalization rule +def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx, normalization): + if not normalization: + return adj_part + + adj_part = adj_part.coalesce() + deg = torch.histc(adj_matrix[0].float(), bins=node_count) + deg = deg.pow(-0.5) + + row_len = adj_part.size(0) + col_len = adj_part.size(1) + + dleft = torch.sparse_coo_tensor([np.arange(0, row_len).tolist(), + np.arange(0, row_len).tolist()], + deg[row_vtx:(row_vtx + row_len)].float(), + size=(row_len, row_len), + requires_grad=False, device=torch.device("cpu")) + + dright = torch.sparse_coo_tensor([np.arange(0, col_len).tolist(), + np.arange(0, col_len).tolist()], + deg[col_vtx:(col_vtx + col_len)].float(), + size=(col_len, col_len), + requires_grad=False, device=torch.device("cpu")) + # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) + ad_ind, ad_val = torch_sparse.spspmm(adj_part._indices(), adj_part._values(), + dright._indices(), dright._values(), + adj_part.size(0), adj_part.size(1), dright.size(1)) + + adj_part_ind, adj_part_val = torch_sparse.spspmm(dleft._indices(), dleft._values(), + ad_ind, ad_val, + dleft.size(0), dleft.size(1), adj_part.size(1)) + + adj_part = torch.sparse_coo_tensor(adj_part_ind, adj_part_val, + size=(adj_part.size(0), adj_part.size(1)), + requires_grad=False, device=torch.device("cpu")) + + return adj_part + +# Split a COO into partitions of size n_per_proc +# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. +def split_coo(adj_matrix, node_count, n_per_proc, dim): + vtx_indices = list(range(0, node_count, n_per_proc)) + vtx_indices.append(node_count) + + am_partitions = [] + for i in range(len(vtx_indices) - 1): + am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] + am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] + am_part[dim] -= vtx_indices[i] + am_partitions.append(am_part) + + return am_partitions, vtx_indices + +def row_normalize(mx): + """Row-normalize sparse matrix""" + rowsum = torch.sparse.sum(mx, 1) + r_inv = torch.float_power(rowsum, -1).flatten() + # r_inv._values = r_inv._values()[torch.isinf(r_inv._values())] = 0. + # r_mat_inv = torch.diag(r_inv._values()) + r_inv_values = torch.cuda.DoubleTensor(r_inv.size(0)).fill_(0) + r_inv_values[r_inv._indices()[0,:]] = r_inv._values() + # r_inv_values = r_inv._values() + r_inv_values[torch.isinf(r_inv_values)] = 0 + r_mat_inv = torch.sparse_coo_tensor([np.arange(0, r_inv.size(0)).tolist(), + np.arange(0, r_inv.size(0)).tolist()], + r_inv_values, + size=(r_inv.size(0), r_inv.size(0))) + # mx = r_mat_inv.mm(mx.float()) + mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values(), + mx._indices(), mx._values(), + r_mat_inv.size(0), r_mat_inv.size(1), mx.size(1), + coalesced=True) + mx = torch.sparse_coo_tensor(indices=mx_indices, values=mx_values.double(), size=(r_mat_inv.size(0), mx.size(1))) + return mx + +def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, \ + normalize, replicate_graph): + node_count = inputs.size(0) + # n_per_proc = math.ceil(float(node_count) / size) + # n_per_proc = math.ceil(float(node_count) / (size / replication)) + n_per_proc = node_count // (size // replication) + + am_partitions = None + am_pbyp = None + + # inputs = inputs.to(torch.device("cpu")) + # adj_matrix = adj_matrix.to(torch.device("cpu")) + # torch.cuda.synchronize() + + rank_c = rank // replication + # Compute the adj_matrix and inputs partitions for this process + # TODO: Maybe I do want grad here. Unsure. + with torch.no_grad(): + # Column partitions + if not replicate_graph: + am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) + else: + am_partitions = None + # # proc_node_count = vtx_indices[rank_c + 1] - vtx_indices[rank_c] + # # am_pbyp, _ = split_coo(am_partitions[rank_c], node_count, n_per_proc, 0) + # # print(f"before", flush=True) + # # for i in range(len(am_pbyp)): + # # if i == size // replication - 1: + # # last_node_count = vtx_indices[i + 1] - vtx_indices[i] + # # am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), + # # size=(last_node_count, proc_node_count), + # # requires_grad=False) + + # # am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + # # vtx_indices[rank_c], normalize) + # # else: + # # am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), + # # size=(n_per_proc, proc_node_count), + # # requires_grad=False) + + # # am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + # # vtx_indices[rank_c], normalize) + + if not replicate_graph: + for i in range(len(am_partitions)): + proc_node_count = vtx_indices[i + 1] - vtx_indices[i] + am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], + torch.ones(am_partitions[i].size(1)).double(), + size=(node_count, proc_node_count), + requires_grad=False) + am_partitions[i] = scale_elements(adj_matrix, am_partitions[i], node_count, 0, vtx_indices[i], \ + normalize) + adj_matrix_loc = am_partitions[rank_c] + else: + adj_matrix_loc = None + + # # input_partitions = torch.split(inputs, math.ceil(float(inputs.size(0)) / (size / replication)), dim=0) + input_partitions = torch.split(inputs, inputs.size(0) // (size // replication), dim=0) + if len(input_partitions) > (size // replication): + input_partitions_fused = [None] * (size // replication) + input_partitions_fused[:-1] = input_partitions[:-2] + input_partitions_fused[-1] = torch.cat(input_partitions[-2:], dim=0) + input_partitions = input_partitions_fused + + inputs_loc = input_partitions[rank_c] + + # print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) + print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}") + return inputs_loc, adj_matrix_loc, am_partitions, input_partitions + +def one5d_partition_mb(rank, size, batches, replication, batch_size): + rank_c = rank // replication + # batch_partitions = torch.split(batches, int(mb_count // (size / replication)), dim=0) + batch_partitions = torch.split(batches, int(batch_size // (size / replication)), dim=1) + return batch_partitions[rank_c] + # batch_partitions = torch.split(batches, int(mb_count // size), dim=0) + # return batch_partitions[rank] + +def main(local_rank, args, trainset, testset, + g_loc_crows, g_loc_cols, g_loc_vals, + g_loc_t_crows, g_loc_t_cols, g_loc_t_vals, node_count, edge_count, + node_features, edge_features, num_features, num_classes): + + torch.set_num_threads(1) + # load and preprocess dataset + # Initialize distributed environment with SLURM + if "SLURM_PROCID" in os.environ.keys(): + os.environ["RANK"] = os.environ["SLURM_PROCID"] + global_rank = int(os.environ["RANK"]) * args.gpu + local_rank + else: + global_rank = local_rank + + if "SLURM_NTASKS" in os.environ.keys(): + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + global_size = int(os.environ["WORLD_SIZE"]) * args.gpu + + os.environ["MASTER_ADDR"] = args.hostname + os.environ["MASTER_PORT"] = "1234" + + print(f"device_count: {torch.cuda.device_count()}") + print(f"hostname: {socket.gethostname()}") + print(f"local_rank: {local_rank}", flush=True) + if not dist.is_initialized(): + print(f"global_rank: {global_rank} local_rank: {local_rank} size: {global_size} initializing process group", flush=True) + dist.init_process_group(backend=args.dist_backend, rank=global_rank, world_size=global_size, timeout=timedelta(minutes=30)) + dist.barrier() + rank = dist.get_rank() + size = dist.get_world_size() + print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}") + torch.cuda.set_device(local_rank) + + if args.wandb and rank == 0: + wandb.init(project="exatrkx") + + device = torch.device(f'cuda:{rank % args.gpu}') + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_vals) + g_loc_t = torch.sparse_csr_tensor(g_loc_t_crows, g_loc_t_cols, g_loc_t_vals) + g_loc = g_loc.to(device) + g_loc_t = g_loc_t.to(device) + # trainset = trainset.to(device) + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + start_inner_timer = torch.cuda.Event(enable_timing=True) + stop_inner_timer = torch.cuda.Event(enable_timing=True) + + s1 = torch.cuda.Stream() + s2 = torch.cuda.Stream() + + path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', args.dataset) + + torch.manual_seed(1) + + proc_row = size // args.replication + rank_row = rank // args.replication + rank_col = rank % args.replication + if rank_row >= (size // args.replication): + return + + # row_groups, col_groups = get_proc_groups(rank, size, args.replication) + row_groups = None + col_groups = None + + # if local_rank > 0: + # model = model.to(device) + model = InteractionGNN(num_features, + args.n_hidden, + num_classes, + args.nb_node_layer, + args.nb_edge_layer, + args.n_graph_iters, + node_features, + edge_features, + args.layernorm, + args.batchnorm, + args.hidden_activation, + args.output_activation, + args.aggr, + args.checkpointing, + args.replication, + device, + impl=args.impl, + dataset=args.dataset) + + print(f"model: {model}") + model = model.to(device) + model = DistributedDataParallel(model, device_ids=[local_rank], output_device=[local_rank], find_unused_parameters=True) + + # use optimizer + print(f"lr: {args.lr}") + optimizer = torch.optim.AdamW( + model.parameters(), + lr=args.lr, + # lr=args.lr * size, + # lr=args.lr * math.sqrt(size), + betas=(0.9, 0.999), + # betas=(1 - size * (1 - 0.9), 1 - size * (1 - 0.999)), + # eps=1e-08 / math.sqrt(size), + eps=1e-08, + amsgrad=True, + weight_decay=0.01 + ) + print(f"optimizer: {optimizer}") + scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, + step_size=15, + gamma=0.9, + ) + + if args.impl == "torch": + local_batch_size = int(args.batch_size // size) + train_loader = ShaDowKHopSampler(trainset, + depth=args.n_layers, + num_neighbors=args.num_neighbors, + batch_size=args.batch_size, + num_workers=1, + pin_memory=True, + shuffle=False, + drop_last=True) + batch_start = rank * (len(train_loader) // size) + batch_stop = min((rank + 1) * (len(train_loader) // size), len(train_loader)) + batch_count = len(train_loader) + if rank == size - 1: + batch_stop = len(train_loader) + print(f"batch_start: {batch_start} batch_stop: {batch_stop}") + print(f"len(train_loader): {len(train_loader)}") + # train_loader = DataLoader(trainset, + # batch_size=1, + # num_workers=32, + # shuffle=False, + # drop_last=True) + print(f"train_loader depth: {args.n_layers} nn: {args.num_neighbors} batch_size: {args.batch_size}") + print(f"len(train_loader): {len(train_loader)}") + # trainset.edge_index = trainset.adj_t.coo() + test_loader = DataLoader(testset, batch_size=1, num_workers=1) + + row_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + if rank == size - 1: + row_n_bulkmb = args.n_bulkmb - row_n_bulkmb * (proc_row - 1) + + rank_n_bulkmb = int(row_n_bulkmb / args.replication) + if rank == size - 1: + rank_n_bulkmb = row_n_bulkmb - rank_n_bulkmb * (args.replication - 1) + + print(f"global_rank: {global_rank} rank_n_bulkmb: {rank_n_bulkmb} bulkmb: {args.n_bulkmb} node_count: {node_count}") + torch.cuda.empty_cache() + torch.cuda.synchronize() + + dur = [] + for epoch in range(args.n_epochs): + model.train() + if (args.warmup != -1) and (epoch < args.warmup): + lr_scale = min(1.0, float(epoch + 1) / args.warmup) + for pg in optimizer.param_groups: + pg["lr"] = lr_scale * args.lr # * size + # current_lr = optimizer.param_groups[0]['lr'] + + print(f"Epoch: {epoch:04}") + if epoch >= 1: + epoch_start = time.time() + + if args.impl == "torch": + # for i, (batch, n_id) in enumerate(train_loader): + + batch_count = args.batchcount + # for i in range(0, batch_count, size): + # batch_ids = torch.arange((i + rank) * args.batch_size, (i + rank + 1) * args.batch_size) + for i in range(0, batch_count): + batches_start = i * args.batch_size + rank * local_batch_size + batches_stop = i * args.batch_size + (rank + 1) * local_batch_size + batch_ids = torch.arange(batches_start, batches_stop) + if epoch >= 1: + start_time(start_timer, timing_arg=True) + batch = train_loader.__collate__(batch_ids) + if epoch >= 1: + model.module.timings["sample"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + + print(f"Rank: {rank:03} Epoch: {epoch:03d} Batch {i:07}") + optimizer.zero_grad() + if epoch >= 1: + start_time(start_timer, timing_arg=True) + batch = batch.to(device) + logits = model(batch, epoch) + loss, pos_loss, neg_loss = model.module.loss_function(logits, batch) + + # print(f"rank: {rank} batch {i} loss: {loss} pos_loss: {pos_loss} neg_loss: {neg_loss}", flush=True) + # if args.wandb and rank == 0: + # wandb.log({'loss': loss.item(), + # 'pos_loss': pos_loss.item(), + # 'neg_loss': neg_loss.item(), + # 'epoch': epoch}) + loss.backward() + if epoch >= 1: + model.module.timings["train"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + + # if epoch >= 1: + # start_time(start_timer, timing_arg=True) + # # for W in model.parameters(): + # # if W.grad is not None: + # # dist.all_reduce(W.grad.data) + # # W.grad.data /= size + # model.ignn_allreduce(model.named_parameters(), size, epoch) + # if epoch >= 1: + # model.module.timings["allreduce"].append(stop_time(start_timer, stop_timer, barrier=True, timing_arg=True)) + optimizer.step() + + elif args.impl == "cagnet": + batches_all = torch.arange(node_count).to(device) + batch_count = args.batchcount + + last_iter = False + for b in range(0, batch_count, args.n_bulkmb): + # for b in range(batch_start, batch_stop, args.n_bulkmb): + print(f"Rank: {rank:03} Batch Bulk: {b:07}") + if b + args.n_bulkmb > batch_count: + last_iter = True + old_nbulkmb = args.n_bulkmb + old_rank_nbulkmb = rank_n_bulkmb + args.n_bulkmb = batch_count - b + rank_n_bulkmb = batch_count - b + batches_start = b * args.batch_size + batches_stop = (b + args.n_bulkmb) * args.batch_size + batches = batches_all[batches_start:batches_stop].view(args.n_bulkmb, args.batch_size) + # batches_loc = torch.arange(b * args.batch_size, (b + args.n_bulkmb) * args.batch_size).view(args.n_bulkmb, args.batch_size).to(device) + batches_loc = one5d_partition_mb(rank, size, batches, 1, args.batch_size) + + batches_indices_rows = torch.arange(batches_loc.size(0), dtype=torch.int32, device=device) + batches_indices_rows = batches_indices_rows.repeat_interleave(batches_loc.size(1)) + # batches_indices_cols = batches_loc.view(-1) + batches_indices_cols = batches_loc.reshape(-1) + batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) + batches_values = torch.cuda.FloatTensor(batches_loc.size(1) * batches_loc.size(0), + device=device).fill_(1.0) + batches_loc = torch.sparse_coo_tensor(batches_indices, batches_values, + (batches_loc.size(0), node_count)) + + node_count = trainset.x.size(0) + if args.n_darts == -1: + avg_degree = int(edge_count / node_count) + args.n_darts = [] + avg_degree = edge_count / node_count + for d in range(args.n_layers): + dart_count = int(avg_degree * args.samp_num / avg_degree) + args.n_darts.append(dart_count) + + if args.replicate_graph: + rep_pass = 1 + else: + rep_pass = args.replication + + if epoch >= 1: + start_time(start_timer, timing_arg=True) + torch.cuda.nvtx.range_push("sample") + local_batch_size = args.batch_size // size + frontiers_bulk, adj_matrices_bulk = shadow_sampler(g_loc_t, + batches_loc, local_batch_size, \ + [args.num_neighbors] * args.n_layers, + args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + rep_pass, rank, size, row_groups, \ + col_groups, args.timing, \ + model.module.timings, \ + args.replicate_graph, epoch) + if epoch >= 1: + torch.cuda.synchronize() + model.module.timings["sample"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + torch.cuda.nvtx.range_pop() # sample + + if epoch >= 1: + start_time(start_timer, timing_arg=True) + + # First iteration before overlap + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("extract") + adj_matrix = adj_matrices_bulk[0].to_sparse_coo() + + frontier_cpu = frontiers_bulk[0].cpu() + edge_ids_cpu = adj_matrix._values().cpu() + + if args.dataset == "physics_ex3": + batch_train = Batch(batch=frontiers_bulk[0], + edge_index=adj_matrices_bulk[0]._indices(), + y=trainset.y[edge_ids_cpu], + z=trainset.z[frontier_cpu], + weights=trainset.weights[edge_ids_cpu], + r=trainset.r[frontier_cpu], + phi=trainset.phi[frontier_cpu], + eta=trainset.eta[frontier_cpu], + truth_map=trainset.truth_map) + elif args.dataset == "ctd": + batch_train = Batch(hit_id=trainset.hit_id[frontier_cpu], + x=trainset.x[frontier_cpu], + y=trainset.y[edge_ids_cpu], + z=trainset.z[frontier_cpu], + edge_index=adj_matrices_bulk[0]._indices(), + truth_map=trainset.truth_map, + weights=trainset.weights[edge_ids_cpu], + r=trainset.r[frontier_cpu], + phi=trainset.phi[frontier_cpu], + eta=trainset.eta[frontier_cpu], + cluster_r_1=trainset.cluster_r_1[frontier_cpu], + cluster_phi_1=trainset.cluster_phi_1[frontier_cpu], + cluster_z_1=trainset.cluster_z_1[frontier_cpu], + cluster_eta_1=trainset.cluster_eta_1[frontier_cpu], + cluster_r_2=trainset.cluster_r_2[frontier_cpu], + cluster_phi_2=trainset.cluster_phi_2[frontier_cpu], + cluster_z_2=trainset.cluster_z_2[frontier_cpu], + cluster_eta_2=trainset.cluster_eta_2[frontier_cpu], + dr=trainset.dr[edge_ids_cpu], + dphi=trainset.dphi[edge_ids_cpu], + dz=trainset.dz[edge_ids_cpu], + deta=trainset.deta[edge_ids_cpu], + phislope=trainset.phislope[edge_ids_cpu], + rphislope=trainset.rphislope[edge_ids_cpu], + ) + if epoch >= 1: + torch.cuda.synchronize() + model.module.timings["extract"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + torch.cuda.nvtx.range_pop() # extract + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("h2d") + optimizer.zero_grad() + batch_train = batch_train.to(device) + if epoch >= 1: + torch.cuda.synchronize() + model.module.timings["h2d"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + torch.cuda.nvtx.range_pop() # h2d + + for i in range(1, args.n_bulkmb): + with torch.cuda.stream(s1): + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("extract") + adj_matrix = adj_matrices_bulk[i].to_sparse_coo() + + frontier_cpu = frontiers_bulk[i].cpu() + edge_ids_cpu = adj_matrix._values().cpu() + + if args.dataset == "physics_ex3": + batch = Batch(batch=frontiers_bulk[i], + edge_index=adj_matrices_bulk[i]._indices(), + y=trainset.y[edge_ids_cpu], + z=trainset.z[frontier_cpu], + weights=trainset.weights[edge_ids_cpu], + r=trainset.r[frontier_cpu], + phi=trainset.phi[frontier_cpu], + eta=trainset.eta[frontier_cpu], + truth_map=trainset.truth_map) + elif args.dataset == "ctd": + batch = Batch(hit_id=trainset.hit_id[frontier_cpu].pin_memory(), + x=trainset.x[frontier_cpu].pin_memory(), + y=trainset.y[edge_ids_cpu].pin_memory(), + z=trainset.z[frontier_cpu].pin_memory(), + edge_index=adj_matrices_bulk[i]._indices(), + truth_map=trainset.truth_map.pin_memory(), + weights=trainset.weights[edge_ids_cpu].pin_memory(), + r=trainset.r[frontier_cpu].pin_memory(), + phi=trainset.phi[frontier_cpu].pin_memory(), + eta=trainset.eta[frontier_cpu].pin_memory(), + cluster_r_1=trainset.cluster_r_1[frontier_cpu].pin_memory(), + cluster_phi_1=trainset.cluster_phi_1[frontier_cpu].pin_memory(), + cluster_z_1=trainset.cluster_z_1[frontier_cpu].pin_memory(), + cluster_eta_1=trainset.cluster_eta_1[frontier_cpu].pin_memory(), + cluster_r_2=trainset.cluster_r_2[frontier_cpu].pin_memory(), + cluster_phi_2=trainset.cluster_phi_2[frontier_cpu].pin_memory(), + cluster_z_2=trainset.cluster_z_2[frontier_cpu].pin_memory(), + cluster_eta_2=trainset.cluster_eta_2[frontier_cpu].pin_memory(), + dr=trainset.dr[edge_ids_cpu].pin_memory(), + dphi=trainset.dphi[edge_ids_cpu].pin_memory(), + dz=trainset.dz[edge_ids_cpu].pin_memory(), + deta=trainset.deta[edge_ids_cpu].pin_memory(), + phislope=trainset.phislope[edge_ids_cpu].pin_memory(), + rphislope=trainset.rphislope[edge_ids_cpu].pin_memory(), + ) + if epoch >= 1: + s1.synchronize() + model.module.timings["extract"].append(stop_time(start_inner_timer, stop_inner_timer, s1, timing_arg=True)) + torch.cuda.nvtx.range_pop() # extract + + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("h2d") + # print(f"rank: {rank} copying batch {b + i}", flush=True) + optimizer.zero_grad() + batch = batch.to(device) + # print(f"rank: {rank} done copying batch {b + i}", flush=True) + if epoch >= 1: + s1.synchronize() + model.module.timings["h2d"].append(stop_time(start_inner_timer, stop_inner_timer, s1, timing_arg=True)) + torch.cuda.nvtx.range_pop() # h2d + with torch.cuda.stream(s2): + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("fwd") + # print(f"rank: {rank} fwd batch {b + i}", flush=True) + logits = model(batch_train, epoch) + # print(f"rank: {rank} done fwd batch {b + i}", flush=True) + if epoch >= 1: + s2.synchronize() + model.module.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer, s2, timing_arg=True)) + torch.cuda.nvtx.range_pop() # fwd + + loss, pos_loss, neg_loss = model.module.loss_function(logits, batch_train) + print(f"Rank: {rank:03} Epoch: {epoch:03d} Batch {(b + i):07}") + if args.wandb and rank == 0: + wandb.log({'loss': loss.item(), + 'pos_loss': pos_loss.item(), + 'neg_loss': neg_loss.item(), + 'epoch': epoch}) + if epoch >= 1: + start_time(start_inner_timer, timing_arg=True) + torch.cuda.nvtx.range_push("bwd") + # print(f"rank: {rank} bwd batch {b + i}", flush=True) + loss.backward() + optimizer.step() + # print(f"rank: {rank} done bwd batch {b + i}", flush=True) + if epoch >= 1: + s2.synchronize() + model.module.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer, s2, timing_arg=True)) + torch.cuda.nvtx.range_pop() # bwd + torch.cuda.synchronize() + batch_train = batch + # if epoch >= 1: + # start_time(start_inner_timer, timing_arg=True) + # # for W in model.parameters(): + # # if W.grad is not None: + # # dist.all_reduce(W.grad) + # # W.grad /= size + # # model.module.ignn_allreduce(model.named_parameters(), size, epoch) + # if epoch >= 1: + # model.module.timings["allreduce"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + + if last_iter: + last_iter = False + args.n_bulkmb = old_nbulkmb + rank_n_bulkmb = old_rank_nbulkmb + if epoch >= 1: + torch.cuda.synchronize() + model.module.timings["train"].append(stop_time(start_timer, stop_timer, timing_arg=True)) + + torch.cuda.synchronize() + dist.barrier() + scheduler.step() + curr_lr = scheduler.get_last_lr() + print(f"Epoch: {epoch} lr: {curr_lr}", flush=True) + if epoch >= 1: + dur.append(time.time() - epoch_start) + print(f"Epoch time: {(np.sum(dur) / epoch):0.03f}", flush=True) + + # if rank == 0 and epoch >= 1: # and (epoch % 5 == 0 or epoch == args.n_epochs - 1): + if epoch >= 1: # and (epoch % 5 == 0 or epoch == args.n_epochs - 1): + if args.timing: + sample_dur = [x / 1000 for x in model.module.timings["sample"]] + train_dur = [x / 1000 for x in model.module.timings["train"]] + extract_dur = [x / 1000 for x in model.module.timings["extract"]] + fwdbwd_dur = [x / 1000 for x in model.module.timings["fwdbwd"]] + fwd_dur = [x / 1000 for x in model.module.timings["fwd"]] + bwd_dur = [x / 1000 for x in model.module.timings["bwd"]] + h2d_dur = [x / 1000 for x in model.module.timings["h2d"]] + allreduce_dur = [x / 1000 for x in model.module.timings["allreduce"]] + allreduce_pt1_dur = [x / 1000 for x in model.module.timings["allreduce-pt1"]] + allreduce_pt2_dur = [x / 1000 for x in model.module.timings["allreduce-pt2"]] + allreduce_pt3_dur = [x / 1000 for x in model.module.timings["allreduce-pt3"]] + allreduce_pt4_dur = [x / 1000 for x in model.module.timings["allreduce-pt4"]] + + shadow_sampling_dur = [x / 1000 for x in model.module.timings["shadow-sampling"]] + shadow_selection_dur = [x / 1000 for x in model.module.timings["shadow-selection"]] + shadow_extraction_dur = [x / 1000 for x in model.module.timings["shadow-extraction"]] + + fwd_prelude_dur = [x / 1000 for x in model.module.timings["fwd-prelude"]] + fwd_encoder_dur = [x / 1000 for x in model.module.timings["fwd-encoder"]] + fwd_graphiters_dur = [x / 1000 for x in model.module.timings["fwd-graphiters"]] + + # print(f"Sample: {(np.sum(sample_dur) / epoch):0.03f} Train: {(np.sum(train_dur) / epoch):0.03f}", flush=True) + # print(f"avg(allreduce_pt2_dur): {sum(allreduce_pt2_dur) / len(allreduce_pt2_dur)}", flush=True) + # print(f"med(allreduce_pt2_dur): {median(allreduce_pt2_dur)}", flush=True) + print(f"Rank: {rank} Sample: {(np.sum(sample_dur) / epoch):0.03f} Train: {(np.sum(train_dur) / epoch):0.03f} AllReduce: {(np.sum(allreduce_dur) / epoch):0.03f}", flush=True) + print(f"Rank: {rank} AllReducePt1: {(np.sum(allreduce_pt1_dur) / epoch):0.03f} AllReducePt2: {(np.sum(allreduce_pt2_dur) / epoch):0.03f}") + print(f"Rank: {rank} AllReducePt3: {(np.sum(allreduce_pt3_dur) / epoch):0.03f} AllReducePt4: {(np.sum(allreduce_pt4_dur) / epoch):0.03f}") + print(f"Rank: {rank} Extract: {(np.sum(extract_dur) / epoch):0.03f} FwdBwd: {(np.sum(fwdbwd_dur) / epoch):0.03f}") + print(f"Rank: {rank} Fwd: {(np.sum(fwd_dur) / epoch):0.03f} Bwd: {(np.sum(bwd_dur) / epoch):0.03f} H2D: {(np.sum(h2d_dur) / epoch):0.03f}") + print(f"Rank: {rank} ShaDow-Sampling: {(np.sum(shadow_sampling_dur) / epoch):0.03f} ShaDow-Selection: {(np.sum(shadow_selection_dur) / epoch):0.03f} ShaDow-Extraction: {(np.sum(shadow_extraction_dur) / epoch):0.03f}", flush=True) + print(f"Rank: {rank} Fwd-Prelude: {(np.sum(fwd_prelude_dur) / epoch):0.03f} Fwd-Encoder: {(np.sum(fwd_encoder_dur) / epoch):0.03f} Fwd-Graphiters: {(np.sum(fwd_graphiters_dur) / epoch):0.03f}", flush=True) + + # torch.cuda.synchronize() + # dist.barrier() + + if not args.skip_validation: + model.eval() + print(f"Evaluating", flush=True) + with torch.no_grad(): + # model = model.cpu() + result, masked_auc = model.module.evaluate(test_loader, epoch, curr_lr[0], \ + np.sum(dur), args.wandb, rank) + # model = model.to(device) + if args.wandb: + if args.dataset == "ctd" and rank == 0: + wandb.log({'signal_effiency': result, + 'time': np.sum(dur)}) + elif rank == 0: + wandb.log({'full_auc': result, + 'time': np.sum(dur)}) + + torch.cuda.synchronize() + dist.barrier() + torch.cuda.synchronize() + + # test_loader = DataLoader(testset, batch_size=1, num_workers=1) + # model.eval() + # with torch.no_grad(): + # if rank == 0: + # print(f"Evaluating", flush=True) + # # model = model.cpu() + # # model = model.cuda() + # result, masked_auc = model.module.evaluate(test_loader, args.n_epochs, + # 0.0, 0.0, False, 0) + # # model = model.to(device) + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='IGNN') + parser.add_argument("--dataset", type=str, default="Cora", + help="dataset to train") + parser.add_argument("--sample-method", type=str, default="shadow", + help="sampling algorithm for training") + parser.add_argument("--dropout", type=float, default=0.5, + help="dropout probability") + parser.add_argument("--gpu", type=int, default=4, + help="gpus per node") + parser.add_argument("--lr", type=float, default=0.0002, + help="learning rate") + parser.add_argument("--n-epochs", type=int, default=200, + help="number of training epochs") + parser.add_argument("--n-hidden", type=int, default=128, + help="number of hidden units") + parser.add_argument("--n-layers", type=int, default=1, + help="number of hidden gcn layers") + parser.add_argument("--batch-size", type=int, default=1024, + help="number of vertices in minibatch") + parser.add_argument("--num-neighbors", type=int, default=1, + help="number of neigbors per vertex of minibatch") + parser.add_argument("--weight-decay", type=float, default=5e-4, + help="Weight for L2 loss") + parser.add_argument("--aggr", type=str, default="mean", + help="Aggregator type: mean/sum") + parser.add_argument('--world-size', default=-1, type=int, + help='number of nodes for distributed training') + parser.add_argument('--rank', default=-1, type=int, + help='node rank for distributed training') + parser.add_argument('--dist-url', default='env://', type=str, + help='url used to set up distributed training') + parser.add_argument('--dist-backend', default='nccl', type=str, + help='distributed backend') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument('--normalize', action="store_true", + help='normalize adjacency matrix') + parser.add_argument('--partitioning', default='ONE5D', type=str, + help='partitioning strategy to use') + parser.add_argument('--replication', default=1, type=int, + help='partitioning strategy to use') + parser.add_argument('--n-bulkmb', default=1, type=int, + help='number of minibatches to sample in bulk') + parser.add_argument('--bulk-batch-fetch', default=1, type=int, + help='number of minibatches to fetch features for in bulk') + parser.add_argument('--n-darts', default=-1, type=int, + help='number of darts to throw per minibatch in LADIES \ + sampling') + parser.add_argument('--semibulk', default=128, type=int, + help='number of batches to column extract from in bulk') + parser.add_argument('--timing', action="store_true", + help='whether to turn on timers') + parser.add_argument('--baseline', action="store_true", + help='whether to avoid col selection for baseline \ + comparison') + parser.add_argument('--replicate-graph', action="store_true", + help='replicate adjacency matrix on each device') + parser.add_argument("--nb-node-layer", type=int, default=2, + help="number of hidden node MLP layers") + parser.add_argument("--nb-edge-layer", type=int, default=2, + help="number of hidden edge MLP layers") + parser.add_argument("--n-graph-iters", type=int, default=8, + help="number of message passing iterations") + parser.add_argument('--batchnorm', action="store_true", + help='use batchnorm in mlps') + parser.add_argument('--layernorm', action="store_true", + help='use batchnorm in mlps') + parser.add_argument('--hidden-activation', choices=["SiLU", "ReLU", "Tanh"], + help='hidden activation in mlps') + parser.add_argument('--output-activation', choices=["SiLU", "ReLU", "Tanh"], + help='output activation in mlps') + parser.add_argument('--impl', default='cagnet', type=str, + help='sampling implementation to use (torch/cagnet)') + parser.add_argument('--wandb', action="store_true", + help='use wandb logging') + parser.add_argument('--checkpointing', action="store_true", + help='use checkpointing') + parser.add_argument("--warmup", type=int, default=5, + help="number of warmup iterations for learning rate") + parser.add_argument("--batchcount", type=int, default=5, + help="number of warmup iterations for learning rate") + parser.add_argument('--skip-validation', action="store_true", + help='skip validation step when expensive') + args = parser.parse_args() + args.samp_num = args.num_neighbors + + print(args) + + if args.dataset == "cora" or args.dataset == "reddit": + if args.dataset == "cora": + dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures()) + elif args.dataset == "reddit": + dataset = Reddit(path) + + data = dataset[0] + data = data.to(device) + # data.x.requires_grad = True + inputs = data.x.to(device) + # inputs.requires_grad = True + data.y = data.y.to(device) + edge_index = data.edge_index + num_features = dataset.num_features + num_classes = dataset.num_classes + adj_matrix = edge_index + elif args.dataset == "physics_ex3": + print(f"Loading coo...", flush=True) + input_dir = "/pscratch/sd/a/alokt/data_dir/Example_3/metric_learning/" + with open(f"gnn_train_{args.dataset}.yaml") as stream: + hparams = yaml.safe_load(stream) + + print(f"hparams: {hparams}", flush=True) + + dataset = GraphDataset(input_dir, "trainset", 80, "fit", hparams) + + print(f"dataset: {dataset}", flush=True) + trainset = [] + for data in dataset: + data_obj = Data(hit_id=data["hit_id"], + x=data["x"], + y=data["y"], + z=data["z"], + r=data["r"], + phi=data["phi"], + eta=data["eta"], + edge_index=data["edge_index"], + truth_map=data["truth_map"], + weights=data["weights"]) + + # data_obj = dataset.preprocess_event(data_obj) + trainset.append(data_obj) + trainset = Batch.from_data_list(trainset) + # trainset = trainset.to(device) + trainset.x.share_memory_() + trainset.y.share_memory_() + trainset.z.share_memory_() + trainset.hit_id.share_memory_() + trainset.edge_index.share_memory_() + trainset.truth_map.share_memory_() + trainset.weights.share_memory_() + trainset.r.share_memory_() + trainset.phi.share_memory_() + trainset.eta.share_memory_() + print(f"trainset: {trainset}", flush=True) + + node_count = torch.max(trainset.edge_index) + 1 + edge_count = trainset.edge_index.size(1) + node_features = ["r", "phi", "z"] + edge_features = [] + + testset = GraphDataset(input_dir, "testset", 10, "fit", hparams) + print(f"testset: {testset}", flush=True) + num_features = len(node_features) + num_classes = 2 + + g_loc_indices = trainset.edge_index + g_loc_values = torch.arange(g_loc_indices.size(1), dtype=torch.int64) + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values) + g_loc = g_loc.to_sparse_csr() + g_loc_t = g_loc.t().to_sparse_csr() + print(f"g_loc: {g_loc}", flush=True) + + elif args.dataset == "ctd": + print(f"Loading coo...", flush=True) + # input_dir = "/global/cfs/cdirs/m4439/CTD2023_results/module_map/" + input_dir = "/global/cfs/cdirs/m4439/REL2024/metric_learning_7_30/" + with open(f"gnn_train_{args.dataset}.yaml") as stream: + hparams = yaml.safe_load(stream) + + print(f"hparams: {hparams}", flush=True) + + # trainset = GraphDataset(input_dir, "trainset", 4053, "fit", hparams) + dataset = GraphDataset(input_dir, "trainset", 80, "fit", hparams) + trainset = [] + for data in dataset: + if data is None: + continue + data_obj = Data(hit_id=data["hit_id"], + x=data["x"], + y=data["y"], + z=data["z"], + edge_index=data["edge_index"], + # edge_index=edge_index, + truth_map=data["truth_map"], + weights=data["weights"], + r=data["r"], + phi=data["phi"], + eta=data["eta"], + cluster_r_1=data["cluster_r_1"], + cluster_phi_1=data["cluster_phi_1"], + cluster_z_1=data["cluster_z_1"], + cluster_eta_1=data["cluster_eta_1"], + cluster_r_2=data["cluster_r_2"], + cluster_phi_2=data["cluster_phi_2"], + cluster_z_2=data["cluster_z_2"], + cluster_eta_2=data["cluster_eta_2"], + dr=data["dr"], + dphi=data["dphi"], + dz=data["dz"], + deta=data["deta"], + phislope=data["phislope"], + rphislope=data["rphislope"], + ) + + trainset.append(data_obj) + trainset = Batch.from_data_list(trainset) + trainset.x.share_memory_() + trainset.y.share_memory_() + trainset.z.share_memory_() + trainset.hit_id.share_memory_() + trainset.edge_index.share_memory_() + trainset.truth_map.share_memory_() + trainset.weights.share_memory_() + trainset.r.share_memory_() + trainset.phi.share_memory_() + trainset.eta.share_memory_() + trainset.cluster_r_1.share_memory_() + trainset.cluster_phi_1.share_memory_() + trainset.cluster_z_1.share_memory_() + trainset.cluster_eta_1.share_memory_() + trainset.cluster_r_2.share_memory_() + trainset.cluster_phi_2.share_memory_() + trainset.cluster_z_2.share_memory_() + trainset.cluster_eta_2.share_memory_() + trainset.dr.share_memory_() + trainset.dphi.share_memory_() + trainset.dz.share_memory_() + trainset.deta.share_memory_() + trainset.phislope.share_memory_() + trainset.rphislope.share_memory_() + trainset.batch.share_memory_() + trainset.ptr.share_memory_() + node_count = torch.max(trainset.edge_index) + 1 + edge_count = trainset.edge_index.size(1) + node_features = ["r", "phi", "z", "eta", "cluster_r_1", "cluster_phi_1", + "cluster_z_1", "cluster_eta_1", "cluster_r_2", + "cluster_phi_2", "cluster_z_2", "cluster_eta_2"] + edge_features = ["dr", "dphi", "dz", "deta", "phislope", "rphislope"] + + g_loc_indices = trainset.edge_index + g_loc_values = torch.arange(g_loc_indices.size(1), dtype=torch.int64) + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values) + g_loc = g_loc.to_sparse_csr() + g_loc_t = g_loc.t().to_sparse_csr() + # trainset = trainset.to(device) + print(f"dataset: {dataset}", flush=True) + print(f"trainset: {trainset}", flush=True) + + # input_dir_test = "/global/cfs/cdirs/m4439/CTD2023_results/module_map/" + input_dir_test = "/global/cfs/cdirs/m4439/REL2024/metric_learning_7_30/" + testset = GraphDataset(input_dir_test, "valset", 10, "test", hparams) + print(f"testset: {testset}", flush=True) + + num_features = len(node_features) + num_classes = 2 + + + mp.set_start_method("spawn", force=True) + mp.spawn(main, + args=(args, trainset, testset, + g_loc.crow_indices(), g_loc.col_indices(), g_loc.values(), + g_loc_t.crow_indices(), g_loc_t.col_indices(), g_loc_t.values(), + node_count, edge_count, node_features, edge_features, + num_features, num_classes), + nprocs=args.gpu, + join=True) + diff --git a/examples/shadow_example.py b/examples/shadow_example.py new file mode 100644 index 0000000..cd8434d --- /dev/null +++ b/examples/shadow_example.py @@ -0,0 +1,85 @@ +import os.path as osp + +import torch +import torch.nn.functional as F + +from torch_geometric.datasets import Planetoid, Flickr +from torch_geometric.loader import ShaDowKHopSampler +from torch_geometric.nn import SAGEConv, global_mean_pool +import torch_geometric.transforms as T + +path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', 'cora') +# dataset = Flickr(path) +dataset = Planetoid(path, "cora", transform=T.NormalizeFeatures()) +data = dataset[0] + +kwargs = {'batch_size': 1024, 'num_workers': 6, 'persistent_workers': True} +train_loader = ShaDowKHopSampler(data, depth=2, num_neighbors=5, + node_idx=data.train_mask, **kwargs) +val_loader = ShaDowKHopSampler(data, depth=2, num_neighbors=5, + node_idx=data.val_mask, **kwargs) +test_loader = ShaDowKHopSampler(data, depth=2, num_neighbors=5, + node_idx=data.test_mask, **kwargs) + + +class GNN(torch.nn.Module): + def __init__(self, in_channels, hidden_channels, out_channels): + super().__init__() + self.conv1 = SAGEConv(in_channels, hidden_channels) + self.conv2 = SAGEConv(hidden_channels, hidden_channels) + self.conv3 = SAGEConv(hidden_channels, hidden_channels) + self.lin = torch.nn.Linear(2 * hidden_channels, out_channels) + + def forward(self, x, edge_index, batch, root_n_id): + x = self.conv1(x, edge_index).relu() + x = F.dropout(x, p=0.3) + x = self.conv2(x, edge_index).relu() + x = F.dropout(x, p=0.3, training=self.training) + x = self.conv3(x, edge_index).relu() + x = F.dropout(x, p=0.3, training=self.training) + + # We merge both central node embeddings and subgraph embeddings: + x = torch.cat([x[root_n_id], global_mean_pool(x, batch)], dim=-1) + + x = self.lin(x) + return x + + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +model = GNN(dataset.num_features, 256, dataset.num_classes).to(device) +optimizer = torch.optim.Adam(model.parameters(), lr=0.001) + + +def train(): + model.train() + total_loss = total_examples = 0 + for data in train_loader: + data = data.to(device) + optimizer.zero_grad() + out = model(data.x, data.edge_index, data.batch, data.root_n_id) + loss = F.cross_entropy(out, data.y) + loss.backward() + optimizer.step() + total_loss += float(loss) * data.num_graphs + total_examples += data.num_graphs + return total_loss / total_examples + + +@torch.no_grad() +def test(loader): + model.eval() + total_correct = total_examples = 0 + for data in loader: + data = data.to(device) + out = model(data.x, data.edge_index, data.batch, data.root_n_id) + total_correct += int((out.argmax(dim=-1) == data.y).sum()) + total_examples += data.num_graphs + return total_correct / total_examples + + +for epoch in range(1, 51): + loss = train() + val_acc = test(val_loader) + test_acc = test(test_loader) + print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, ', + f'Val: {val_acc:.4f} Test: {test_acc:.4f}') diff --git a/examples/test_spgemm.py b/examples/test_spgemm.py new file mode 100644 index 0000000..bec6bd3 --- /dev/null +++ b/examples/test_spgemm.py @@ -0,0 +1,6 @@ +import torch +from sparse_coo_tensor_cpp import spgemm_gpu + +x = torch.sparse_coo_tensor(indices=torch.cuda.LongTensor([[0, 1, 1], [1, 0, 0]]), values=torch.cuda.DoubleTensor([2, 2, 3])) +y = torch.sparse_coo_tensor(indices=torch.cuda.LongTensor([[0, 1], [1, 0]]), values=torch.cuda.DoubleTensor([3, 3])) +print(f"{spgemm_gpu(x._indices()[0,:].int(), x._indices()[1,:].int(), x._values().float(), y._indices()[0,:].int(), y._indices()[1,:].int(), y._values().float(), 2, 2, 2)}") diff --git a/gcn.py b/gcn.py deleted file mode 100644 index 32667b1..0000000 --- a/gcn.py +++ /dev/null @@ -1,126 +0,0 @@ -import os.path as osp -import argparse - -import torch -from torch.nn import Parameter -import torch.nn.functional as F -from torch_geometric.datasets import Planetoid, PPI, Reddit -import torch_geometric.transforms as T -from torch_geometric.nn import GCNConv # noqa - -import time - -parser = argparse.ArgumentParser() -parser.add_argument('--use_gdc', action='store_true', - help='Use GDC preprocessing.') -args = parser.parse_args() - -dataset = 'Cora' -path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset) -dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures()) -# dataset = PPI(path, 'train', T.NormalizeFeatures()) -# dataset = Reddit(path, T.NormalizeFeatures()) -# dataset = Yelp(path, T.NormalizeFeatures()) -data = dataset[0] - -seed = 0 - -if args.use_gdc: - gdc = T.GDC(self_loop_weight=1, normalization_in='sym', - normalization_out='col', - diffusion_kwargs=dict(method='ppr', alpha=0.05), - sparsification_kwargs=dict(method='topk', k=128, - dim=0), exact=True) - data = gdc(data) - - -class Net(torch.nn.Module): - def __init__(self): - super(Net, self).__init__() - self.conv1 = GCNConv(dataset.num_features, 16, cached=True, normalize=False, bias=False) - self.conv2 = GCNConv(16, dataset.num_classes, cached=True, normalize=False, bias=False) - - self.conv1.node_dim = 0 - self.conv2.node_dim = 0 - - with torch.no_grad(): - self.conv1.weight = Parameter(weight1) - self.conv2.weight = Parameter(weight2) - # self.conv1 = ChebConv(data.num_features, 16, K=2) - # self.conv2 = ChebConv(16, data.num_features, K=2) - - def forward(self): - x, edge_index = data.x, data.edge_index - # x = F.relu(self.conv1(x, edge_index)) - # x = F.dropout(x, training=self.training) - # x = self.conv2(x, edge_index) - # return F.log_softmax(x, dim=1) - x = self.conv1(x, edge_index) - # x = F.relu(x) - x = self.conv2(x, edge_index) - # return F.log_softmax(x, dim=1) - return x - - -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - -torch.manual_seed(seed) -weight1 = torch.rand(dataset.num_features, 16) -weight1 = weight1.to(device) - -weight2 = torch.rand(16, dataset.num_classes) -weight2 = weight2.to(device) - -data.y = data.y.type(torch.LongTensor) -model, data = Net().to(device), data.to(device) - -optimizer = torch.optim.Adam(model.parameters(), lr=0.01) - -def train(): - model.train() - optimizer.zero_grad() - outputs = model() - - # Note: bool type removes warnings, unsure of perf penalty - F.nll_loss(outputs[data.train_mask.bool()], data.y[data.train_mask.bool()]).backward() - # F.nll_loss(outputs, torch.max(data.y, 1)[1]).backward() - - for W in model.parameters(): - if W.grad is not None: - print(W.grad) - - optimizer.step() - return outputs - -def test(outputs): - model.eval() - logits, accs = outputs, [] - for _, mask in data('train_mask', 'val_mask', 'test_mask'): - pred = logits[mask].max(1)[1] - acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - accs.append(acc) - return accs - -def main(): - best_val_acc = test_acc = 0 - outputs = None - - tstart = time.time() - - # for epoch in range(1, 101): - for epoch in range(1): - outputs = train() - train_acc, val_acc, tmp_test_acc = test(outputs) - if val_acc > best_val_acc: - best_val_acc = val_acc - test_acc = tmp_test_acc - log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}' - print(log.format(epoch, train_acc, best_val_acc, test_acc)) - - tstop = time.time() - print("Time: " + str(tstop - tstart)) - - return outputs - -if __name__=='__main__': - print(main()) diff --git a/gcn_distr.py b/gcn_distr.py deleted file mode 100644 index 7f68144..0000000 --- a/gcn_distr.py +++ /dev/null @@ -1,838 +0,0 @@ -import os -import os.path as osp -import argparse - -import math - -import torch -import torch.distributed as dist - -from torch_geometric.data import Data, Dataset -from torch_geometric.datasets import Planetoid, PPI -from reddit import Reddit -from torch_geometric.nn import GCNConv, ChebConv # noqa -from torch_geometric.utils import add_remaining_self_loops, to_dense_adj, dense_to_sparse, to_scipy_sparse_matrix -import torch_geometric.transforms as T - -import torch.multiprocessing as mp - -from torch.multiprocessing import Manager, Process - -from torch.nn import Parameter -import torch.nn.functional as F - -from torch_scatter import scatter_add -import torch_sparse - -from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu - -import socket -import statistics -import time -import numpy as np - -# comp_time = 0.0 -# comm_time = 0.0 -# scomp_time = 0.0 -# dcomp_time = 0.0 -# bcast_comm_time = 0.0 -# bcast_words = 0 -# op1_comm_time = 0.0 -# op2_comm_time = 0.0 -total_time = dict() -comp_time = dict() -comm_time = dict() -scomp_time = dict() -dcomp_time = dict() -bcast_comm_time = dict() -barrier_time = dict() -barrier_subset_time = dict() -op1_comm_time = dict() -op2_comm_time = dict() - -epochs = 0 -graphname = "" -mid_layer = 0 -timing = True -normalization = False -activations = False -accuracy = False -device = None -acc_per_rank = 0 -run_count = 0 -run = 0 -download = False - -def start_time(group, rank, subset=False, src=None): - global barrier_time - global barrier_subset_time - - if not timing: - return 0.0 - - barrier_tstart = time.time() - - dist.barrier(group) - barrier_tstop = time.time() - barrier_time[run][rank] += barrier_tstop - barrier_tstart - if subset: - barrier_subset_time[run][rank] += barrier_tstop - barrier_tstart - - tstart = 0.0 - tstart = time.time() - return tstart - -def stop_time(group, rank, tstart): - if not timing: - return 0.0 - dist.barrier(group) - devid = rank_to_devid(rank, acc_per_rank) - device = torch.device('cuda:{}'.format(devid)) - # torch.cuda.synchronize(device=device) - - tstop = time.time() - return tstop - tstart - -def normalize(adj_matrix): - adj_matrix = adj_matrix + torch.eye(adj_matrix.size(0)) - d = torch.sum(adj_matrix, dim=1) - d = torch.rsqrt(d) - d = torch.diag(d) - return torch.mm(d, torch.mm(adj_matrix, d)) - -def block_row(adj_matrix, am_partitions, inputs, weight, rank, size): - n_per_proc = math.ceil(float(adj_matrix.size(1)) / size) - # n_per_proc = int(adj_matrix.size(1) / size) - # am_partitions = list(torch.split(adj_matrix, n_per_proc, dim=1)) - - z_loc = torch.cuda.FloatTensor(n_per_proc, inputs.size(1)).fill_(0) - # z_loc = torch.zeros(adj_matrix.size(0), inputs.size(1)) - - inputs_recv = torch.zeros(inputs.size()) - - part_id = rank % size - - z_loc += torch.mm(am_partitions[part_id].t(), inputs) - - for i in range(1, size): - part_id = (rank + i) % size - - inputs_recv = torch.zeros(am_partitions[part_id].size(0), inputs.size(1)) - - src = (rank + 1) % size - dst = rank - 1 - if dst < 0: - dst = size - 1 - - if rank == 0: - dist.send(tensor=inputs, dst=dst) - dist.recv(tensor=inputs_recv, src=src) - else: - dist.recv(tensor=inputs_recv, src=src) - dist.send(tensor=inputs, dst=dst) - - inputs = inputs_recv.clone() - - # z_loc += torch.mm(am_partitions[part_id], inputs) - z_loc += torch.mm(am_partitions[part_id].t(), inputs) - - - # z_loc = torch.mm(z_loc, weight) - return z_loc - -def outer_product(adj_matrix, grad_output, rank, size, group): - global comm_time - global comp_time - global dcomp_time - global op1_comm_time - global run - - - n_per_proc = math.ceil(float(adj_matrix.size(0)) / size) - - tstart_comp = start_time(group, rank) - - # A * G^l - ag = torch.mm(adj_matrix, grad_output) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - tstart_comm = start_time(group, rank) - - # reduction on A * G^l low-rank matrices - dist.all_reduce(ag, op=dist.reduce_op.SUM, group=group) - - dur = stop_time(group, rank, tstart_comm) - comm_time[run][rank] += dur - op1_comm_time[run][rank] += dur - - # partition A * G^l by block rows and get block row for this process - # TODO: this might not be space-efficient - red_partitions = list(torch.split(ag, n_per_proc, dim=0)) - grad_input = red_partitions[rank] - - return grad_input - -def outer_product2(inputs, ag, rank, size, group): - global comm_time - global comp_time - global dcomp_time - global op2_comm_time - global run - - tstart_comp = start_time(group, rank) - # (H^(l-1))^T * (A * G^l) - grad_weight = torch.mm(inputs, ag) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - tstart_comm = start_time(group, rank) - # reduction on grad_weight low-rank matrices - dist.all_reduce(grad_weight, op=dist.reduce_op.SUM, group=group) - - dur = stop_time(group, rank, tstart_comm) - comm_time[run][rank] += dur - op2_comm_time[run][rank] += dur - - return grad_weight - -def broad_func(node_count, am_partitions, inputs, rank, size, group): - global device - global comm_time - global comp_time - global scomp_time - global bcast_comm_time - global run - - # n_per_proc = math.ceil(float(adj_matrix.size(1)) / size) - n_per_proc = math.ceil(float(node_count) / size) - - # z_loc = torch.cuda.FloatTensor(adj_matrix.size(0), inputs.size(1), device=device).fill_(0) - z_loc = torch.cuda.FloatTensor(am_partitions[0].size(0), inputs.size(1), device=device).fill_(0) - # z_loc = torch.zeros(adj_matrix.size(0), inputs.size(1)) - - inputs_recv = torch.cuda.FloatTensor(n_per_proc, inputs.size(1), device=device).fill_(0) - # inputs_recv = torch.zeros(n_per_proc, inputs.size(1)) - - for i in range(size): - if i == rank: - inputs_recv = inputs.clone() - elif i == size - 1: - inputs_recv = torch.cuda.FloatTensor(am_partitions[i].size(1), inputs.size(1), device=device).fill_(0) - # inputs_recv = torch.zeros(list(am_partitions[i].t().size())[1], inputs.size(1)) - - tstart_comm = start_time(group, rank) - - dist.broadcast(inputs_recv, src=i, group=group) - - dur = stop_time(group, rank, tstart_comm) - comm_time[run][rank] += dur - bcast_comm_time[run][rank] += dur - - tstart_comp = start_time(group, rank) - - spmm_gpu(am_partitions[i].indices()[0].int(), am_partitions[i].indices()[1].int(), - am_partitions[i].values(), am_partitions[i].size(0), - am_partitions[i].size(1), inputs_recv, z_loc) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - scomp_time[run][rank] += dur - - return z_loc - -class GCNFunc(torch.autograd.Function): - @staticmethod - def forward(ctx, inputs, weight, adj_matrix, am_partitions, rank, size, group, func): - global comm_time - global comp_time - global dcomp_time - global run - - # inputs: H - # adj_matrix: A - # weight: W - # func: sigma - - # adj_matrix = adj_matrix.to_dense() - ctx.save_for_backward(inputs, weight, adj_matrix) - ctx.am_partitions = am_partitions - ctx.rank = rank - ctx.size = size - ctx.group = group - - ctx.func = func - - # z = block_row(adj_matrix.t(), am_partitions, inputs, weight, rank, size) - z = broad_func(adj_matrix.size(0), am_partitions, inputs, rank, size, group) - - tstart_comp = start_time(group, rank) - - z = torch.mm(z, weight) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - z.requires_grad = True - ctx.z = z - - if activations: - if func is F.log_softmax: - h = func(z, dim=1) - elif func is F.relu: - h = func(z) - else: - h = z - - return h - else: - return z - - @staticmethod - def backward(ctx, grad_output): - global comm_time - global comp_time - global dcomp_time - global run - - inputs, weight, adj_matrix = ctx.saved_tensors - am_partitions = ctx.am_partitions - rank = ctx.rank - size = ctx.size - group = ctx.group - - func = ctx.func - z = ctx.z - - if activations: - with torch.set_grad_enabled(True): - if func is F.log_softmax: - func_eval = func(z, dim=1) - elif func is F.relu: - func_eval = func(z) - else: - func_eval = z - - sigmap = torch.autograd.grad(outputs=func_eval, inputs=z, grad_outputs=grad_output)[0] - grad_output = sigmap - - - # First backprop equation - ag = broad_func(adj_matrix.size(0), am_partitions, grad_output, rank, size, group) - - tstart_comp = start_time(group, rank) - - grad_input = torch.mm(ag, weight.t()) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - # Second backprop equation (reuses the A * G^l computation) - grad_weight = outer_product2(inputs.t(), ag, rank, size, group) - - return grad_input, grad_weight, None, None, None, None, None, None - -def train(inputs, weight1, weight2, adj_matrix, am_partitions, optimizer, data, rank, size, group): - outputs = GCNFunc.apply(inputs, weight1, adj_matrix, am_partitions, rank, size, group, F.relu) - outputs = GCNFunc.apply(outputs, weight2, adj_matrix, am_partitions, rank, size, group, F.log_softmax) - - optimizer.zero_grad() - - node_count = adj_matrix.size(0) - n_per_proc = int(math.ceil(float(node_count) / size)) - rank_train_mask = torch.split(data.train_mask.bool(), n_per_proc, dim=0)[rank] - datay_rank = torch.split(data.y, n_per_proc, dim=0)[rank] - - # Note: bool type removes warnings, unsure of perf penalty - # loss = F.nll_loss(outputs[data.train_mask.bool()], data.y[data.train_mask.bool()]) - if list(datay_rank[rank_train_mask].size())[0] > 0: - # if datay_rank.size(0) > 0: - loss = F.nll_loss(outputs[rank_train_mask], datay_rank[rank_train_mask]) - # loss = F.nll_loss(outputs, torch.max(datay_rank, 1)[1]) - loss.backward() - else: - fake_loss = (outputs * torch.cuda.FloatTensor(outputs.size(), device=device).fill_(0)).sum() - # fake_loss = (outputs * torch.zeros(outputs.size())).sum() - fake_loss.backward() - - optimizer.step() - - return outputs - -def test(outputs, data, vertex_count, rank): - logits, accs = outputs, [] - for _, mask in data('train_mask', 'val_mask', 'test_mask'): - pred = logits[mask].max(1)[1] - acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - accs.append(acc) - - if len(accs) == 1: - accs.append(0) - accs.append(0) - - return accs - # logits, accs = outputs, [] - # datay_rank = torch.split(data.y, vertex_count)[rank] - # for _, mask in data('train_mask', 'val_mask', 'test_mask'): - # mask_rank = torch.split(mask, vertex_count)[rank] - # count = mask_rank.nonzero().size(0) - # if count > 0: - # pred = logits[mask_rank].max(1)[1] - # acc = pred.eq(datay_rank[mask_rank]).sum().item() / mask_rank.sum().item() - # # pred = logits[mask].max(1)[1] - # # acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - # else: - # acc = -1 - # accs.append(acc) - # return accs - - -# Split a COO into partitions of size n_per_proc -# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. -def split_coo(adj_matrix, node_count, n_per_proc, dim): - vtx_indices = list(range(0, node_count, n_per_proc)) - vtx_indices.append(node_count) - - am_partitions = [] - for i in range(len(vtx_indices) - 1): - am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] - am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] - am_part[dim] -= vtx_indices[i] - am_partitions.append(am_part) - - return am_partitions, vtx_indices - -# Normalize all elements according to KW's normalization rule -def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx): - if not normalization: - return adj_part - - # Scale each edge (u, v) by 1 / (sqrt(u) * sqrt(v)) - # indices = adj_part._indices() - # values = adj_part._values() - - # deg_map = dict() - # for i in range(adj_part._nnz()): - # u = indices[0][i] + row_vtx - # v = indices[1][i] + col_vtx - - # if u.item() in deg_map: - # degu = deg_map[u.item()] - # else: - # degu = (adj_matrix[0] == u).sum().item() - # deg_map[u.item()] = degu - - # if v.item() in deg_map: - # degv = deg_map[v.item()] - # else: - # degv = (adj_matrix[0] == v).sum().item() - # deg_map[v.item()] = degv - - # values[i] = values[i] / (math.sqrt(degu) * math.sqrt(degv)) - - adj_part = adj_part.coalesce() - deg = torch.histc(adj_matrix[0].double(), bins=node_count) - deg = deg.pow(-0.5) - - row_len = adj_part.size(0) - col_len = adj_part.size(1) - - dleft = torch.sparse_coo_tensor([np.arange(0, row_len).tolist(), - np.arange(0, row_len).tolist()], - deg[row_vtx:(row_vtx + row_len)].float(), - size=(row_len, row_len), - requires_grad=False, device=torch.device("cpu")) - - dright = torch.sparse_coo_tensor([np.arange(0, col_len).tolist(), - np.arange(0, col_len).tolist()], - deg[col_vtx:(col_vtx + col_len)].float(), - size=(col_len, col_len), - requires_grad=False, device=torch.device("cpu")) - # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) - ad_ind, ad_val = torch_sparse.spspmm(adj_part._indices(), adj_part._values(), - dright._indices(), dright._values(), - adj_part.size(0), adj_part.size(1), dright.size(1)) - - adj_part_ind, adj_part_val = torch_sparse.spspmm(dleft._indices(), dleft._values(), - ad_ind, ad_val, - dleft.size(0), dleft.size(1), adj_part.size(1)) - - adj_part = torch.sparse_coo_tensor(adj_part_ind, adj_part_val, - size=(adj_part.size(0), adj_part.size(1)), - requires_grad=False, device=torch.device("cpu")) - - return adj_part - -def oned_partition(rank, size, inputs, adj_matrix, data, features, classes, device): - node_count = inputs.size(0) - n_per_proc = math.ceil(float(node_count) / size) - - am_partitions = None - am_pbyp = None - - inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) - - # Compute the adj_matrix and inputs partitions for this process - # TODO: Maybe I do want grad here. Unsure. - with torch.no_grad(): - # Column partitions - am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) - - proc_node_count = vtx_indices[rank + 1] - vtx_indices[rank] - am_pbyp, _ = split_coo(am_partitions[rank], node_count, n_per_proc, 0) - for i in range(len(am_pbyp)): - if i == size - 1: - last_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(last_node_count, proc_node_count), - requires_grad=False) - - am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank]) - else: - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(n_per_proc, proc_node_count), - requires_grad=False) - - am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank]) - - for i in range(len(am_partitions)): - proc_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], - torch.ones(am_partitions[i].size(1)), - size=(node_count, proc_node_count), - requires_grad=False) - am_partitions[i] = scale_elements(adj_matrix, am_partitions[i], node_count, 0, vtx_indices[i]) - - input_partitions = torch.split(inputs, math.ceil(float(inputs.size(0)) / size), dim=0) - - adj_matrix_loc = am_partitions[rank] - inputs_loc = input_partitions[rank] - - print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) - print(f"rank: {rank} inputs.size: {inputs.size()}", flush=True) - return inputs_loc, adj_matrix_loc, am_pbyp - -def run(rank, size, inputs, adj_matrix, data, features, classes, device): - global epochs - global mid_layer - global run - global timing - - best_val_acc = test_acc = 0 - outputs = None - group = dist.new_group(list(range(size))) - - if rank >= size: - return - - # adj_matrix_loc = torch.rand(node_count, n_per_proc) - # inputs_loc = torch.rand(n_per_proc, inputs.size(1)) - - inputs_loc, adj_matrix_loc, am_pbyp = oned_partition(rank, size, inputs, adj_matrix, data, - features, classes, device) - - inputs_loc = inputs_loc.to(device) - adj_matrix_loc = adj_matrix_loc.to(device) - for i in range(len(am_pbyp)): - am_pbyp[i] = am_pbyp[i].t().coalesce().to(device) - - for i in range(run_count): - run = i - torch.manual_seed(0) - weight1_nonleaf = torch.rand(features, mid_layer, requires_grad=True) - weight1_nonleaf = weight1_nonleaf.to(device) - weight1_nonleaf.retain_grad() - - weight2_nonleaf = torch.rand(mid_layer, classes, requires_grad=True) - weight2_nonleaf = weight2_nonleaf.to(device) - weight2_nonleaf.retain_grad() - - weight1 = Parameter(weight1_nonleaf) - weight2 = Parameter(weight2_nonleaf) - - optimizer = torch.optim.Adam([weight1, weight2], lr=0.01) - dist.barrier(group) - - tstart = 0.0 - tstop = 0.0 - - total_time[i] = dict() - comm_time[i] = dict() - comp_time[i] = dict() - scomp_time[i] = dict() - dcomp_time[i] = dict() - bcast_comm_time[i] = dict() - barrier_time[i] = dict() - barrier_subset_time[i] = dict() - op1_comm_time[i] = dict() - op2_comm_time[i] = dict() - - total_time[i][rank] = 0.0 - comm_time[i][rank] = 0.0 - comp_time[i][rank] = 0.0 - scomp_time[i][rank] = 0.0 - dcomp_time[i][rank] = 0.0 - bcast_comm_time[i][rank] = 0.0 - barrier_time[i][rank] = 0.0 - barrier_subset_time[i][rank] = 0.0 - op1_comm_time[i][rank] = 0.0 - op2_comm_time[i][rank] = 0.0 - - timing_on = timing == True - timing = False - outputs = train(inputs_loc, weight1, weight2, adj_matrix_loc, am_pbyp, optimizer, data, - rank, size, group) - if timing_on: - timing = True - - dist.barrier(group) - tstart = time.time() - - # for epoch in range(1, 201): - print(f"Starting training... rank {rank} run {i}", flush=True) - for epoch in range(1, epochs): - outputs = train(inputs_loc, weight1, weight2, adj_matrix_loc, am_pbyp, optimizer, data, - rank, size, group) - print("Epoch: {:03d}".format(epoch), flush=True) - - # dist.barrier(group) - tstop = time.time() - total_time[i][rank] = tstop - tstart - - # Get median runtime according to rank0 and print that run's breakdown - dist.barrier(group) - if rank == 0: - total_times_r0 = [] - for i in range(run_count): - total_times_r0.append(total_time[i][0]) - - print(f"total_times_r0: {total_times_r0}") - median_run_time = statistics.median(total_times_r0) - median_idx = total_times_r0.index(median_run_time) - median_idx = torch.cuda.LongTensor([median_idx]) - else: - median_idx = torch.cuda.LongTensor([0]) - - dist.broadcast(median_idx, src=0, group=group) - median_idx = median_idx.item() - print(f"rank: {rank} median_run: {median_idx}") - print(f"rank: {rank} total_time: {total_time[median_idx][rank]}") - print(f"rank: {rank} comm_time: {comm_time[median_idx][rank]}") - print(f"rank: {rank} comp_time: {comp_time[median_idx][rank]}") - print(f"rank: {rank} scomp_time: {scomp_time[median_idx][rank]}") - print(f"rank: {rank} dcomp_time: {dcomp_time[median_idx][rank]}") - print(f"rank: {rank} bcast_comm_time: {bcast_comm_time[median_idx][rank]}") - print(f"rank: {rank} barrier_time: {barrier_time[median_idx][rank]}") - print(f"rank: {rank} barrier_subset_time: {barrier_subset_time[median_idx][rank]}") - print(f"rank: {rank} op1_comm_time: {op1_comm_time[median_idx][rank]}") - print(f"rank: {rank} op2_comm_time: {op2_comm_time[median_idx][rank]}") - print(f"rank: {rank} {outputs}") - - - if accuracy: - # All-gather outputs to test accuracy - output_parts = [] - n_per_proc = math.ceil(float(inputs.size(0)) / size) - # print(f"rows: {am_pbyp[-1].size(0)} cols: {classes}", flush=True) - for i in range(size): - output_parts.append(torch.cuda.FloatTensor(n_per_proc, classes, device=device).fill_(0)) - - if outputs.size(0) != n_per_proc: - pad_row = n_per_proc - outputs.size(0) - outputs = torch.cat((outputs, torch.cuda.FloatTensor(pad_row, classes, device=device)), dim=0) - - dist.all_gather(output_parts, outputs) - output_parts[rank] = outputs - - padding = inputs.size(0) - n_per_proc * (size - 1) - output_parts[size - 1] = output_parts[size - 1][:padding,:] - - outputs = torch.cat(output_parts, dim=0) - - train_acc, val_acc, tmp_test_acc = test(outputs, data, am_pbyp[0].size(1), rank) - if val_acc > best_val_acc: - best_val_acc = val_acc - test_acc = tmp_test_acc - log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}' - - print(log.format(900, train_acc, best_val_acc, test_acc)) - - return outputs - -def rank_to_devid(rank, acc_per_rank): - return rank % acc_per_rank - -def init_process(rank, size, inputs, adj_matrix, data, features, classes, device, outputs, fn): - run_outputs = fn(rank, size, inputs, adj_matrix, data, features, classes, device) - if outputs is not None: - outputs[rank] = run_outputs.detach() - -def main(): - global device - global graphname - - print(socket.gethostname()) - seed = 0 - - if not download: - mp.set_start_method('spawn', force=True) - outputs = None - if "OMPI_COMM_WORLD_RANK" in os.environ.keys(): - os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] - - # Initialize distributed environment with SLURM - if "SLURM_PROCID" in os.environ.keys(): - os.environ["RANK"] = os.environ["SLURM_PROCID"] - - if "SLURM_NTASKS" in os.environ.keys(): - os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] - - if "MASTER_ADDR" not in os.environ.keys(): - os.environ["MASTER_ADDR"] = "127.0.0.1" - - os.environ["MASTER_PORT"] = "1234" - dist.init_process_group(backend='nccl') - rank = dist.get_rank() - size = dist.get_world_size() - print("Processes: " + str(size)) - - # device = torch.device('cpu') - devid = rank_to_devid(rank, acc_per_rank) - device = torch.device('cuda:{}'.format(devid)) - print(f"device: {device}") - torch.cuda.set_device(device) - curr_devid = torch.cuda.current_device() - # print(f"curr_devid: {curr_devid}", flush=True) - devcount = torch.cuda.device_count() - - if graphname == "Cora": - path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures()) - data = dataset[0] - data = data.to(device) - data.x.requires_grad = True - inputs = data.x.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - edge_index = data.edge_index - num_features = dataset.num_features - num_classes = dataset.num_classes - elif graphname == "Reddit": - path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - dataset = Reddit(path, T.NormalizeFeatures()) - data = dataset[0] - data = data.to(device) - data.x.requires_grad = True - inputs = data.x.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - edge_index = data.edge_index - num_features = dataset.num_features - num_classes = dataset.num_classes - elif graphname == 'Amazon': - # path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - # edge_index = torch.load(path + "/processed/amazon_graph.pt") - # edge_index = torch.load("/gpfs/alpine/bif115/scratch/alokt/Amazon/processed/amazon_graph_jsongz.pt") - # edge_index = edge_index.t_() - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/Amazon/processed/data.pt") - print(f"Done loading coo", flush=True) - # n = 9430088 - n = 14249639 - num_features = 300 - num_classes = 24 - # mid_layer = 24 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1).long() - data.train_mask = torch.ones(n).long() - # edge_index = edge_index.to(device) - print(f"edge_index.size: {edge_index.size()}", flush=True) - print(f"edge_index: {edge_index}", flush=True) - data = data.to(device) - # inputs = inputs.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - elif graphname == 'subgraph3': - # path = "/gpfs/alpine/bif115/scratch/alokt/HipMCL/" - # print(f"Loading coo...", flush=True) - # edge_index = torch.load(path + "/processed/subgraph3_graph.pt") - # print(f"Done loading coo", flush=True) - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/subgraph3/processed/data.pt") - print(f"Done loading coo", flush=True) - n = 8745542 - num_features = 128 - # mid_layer = 512 - # mid_layer = 64 - num_classes = 256 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1).long() - data.train_mask = torch.ones(n).long() - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - - if download: - exit() - - if normalization: - adj_matrix, _ = add_remaining_self_loops(edge_index, num_nodes=inputs.size(0)) - else: - adj_matrix = edge_index - - - init_process(rank, size, inputs, adj_matrix, data, num_features, num_classes, device, outputs, - run) - - if outputs is not None: - return outputs[0] - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument("--local_rank", type=int) - parser.add_argument("--accperrank", type=int) - parser.add_argument("--epochs", type=int) - parser.add_argument("--graphname", type=str) - parser.add_argument("--timing", type=str) - parser.add_argument("--midlayer", type=int) - parser.add_argument("--runcount", type=int) - parser.add_argument("--normalization", type=str) - parser.add_argument("--activations", type=str) - parser.add_argument("--accuracy", type=str) - parser.add_argument("--download", type=bool) - - args = parser.parse_args() - print(args) - - acc_per_rank = args.accperrank - - epochs = args.epochs - graphname = args.graphname - timing = args.timing == "True" - mid_layer = args.midlayer - run_count = args.runcount - normalization = args.normalization == "True" - activations = args.activations == "True" - accuracy = args.accuracy == "True" - download = args.download - - if not download: - if (epochs is None) or (graphname is None) or (timing is None) or (mid_layer is None) or (run_count is None): - print(f"Error: missing argument {epochs} {graphname} {timing} {mid_layer} {run_count}") - exit() - - print(f"Arguments: epochs: {epochs} graph: {graphname} timing: {timing} mid: {mid_layer} norm: {normalization} act: {activations} acc: {accuracy} runs: {run_count}") - - print(main()) diff --git a/gcn_distr_15d.py b/gcn_distr_15d.py deleted file mode 100644 index 587ffb6..0000000 --- a/gcn_distr_15d.py +++ /dev/null @@ -1,874 +0,0 @@ -import os -import os.path as osp -import argparse - -import math - -import torch -import torch.distributed as dist - -from torch_geometric.data import Data, Dataset -from torch_geometric.datasets import Planetoid, PPI -from reddit import Reddit -from torch_geometric.nn import GCNConv, ChebConv # noqa -from torch_geometric.utils import add_remaining_self_loops, to_dense_adj, dense_to_sparse, to_scipy_sparse_matrix -import torch_geometric.transforms as T - -import torch.multiprocessing as mp - -from torch.multiprocessing import Manager, Process - -from torch.nn import Parameter -import torch.nn.functional as F - -from torch_scatter import scatter_add -import torch_sparse - -from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu - -import socket -import statistics -import time -import numpy as np - -# comp_time = 0.0 -# comm_time = 0.0 -# scomp_time = 0.0 -# dcomp_time = 0.0 -# bcast_comm_time = 0.0 -# bcast_words = 0 -# op1_comm_time = 0.0 -# op2_comm_time = 0.0 -total_time = dict() -comp_time = dict() -comm_time = dict() -scomp_time = dict() -dcomp_time = dict() -bcast_comm_time = dict() -bcast_words = dict() -reduce_comm_time = dict() -op_comm_time = dict() -barrier_time = dict() - -epochs = 0 -graphname = "" -mid_layer = 0 -timing = True -normalization = False -activations = False -accuracy = False -device = None -acc_per_rank = 0 -run_count = 0 -run = 0 -replication = 0 -download = False - -def start_time(group, rank, subset=False, src=None): - global barrier_time - global run - - if not timing: - return 0.0 - if group is not None: - barrier_tstart = time.time() - # dist.barrier(group) - torch.cuda.synchronize(device=device) - barrier_tstop = time.time() - barrier_time[run][rank] += barrier_tstop - barrier_tstart - tstart = 0.0 - if rank == 0: - tstart = time.time() - return tstart - -def stop_time(group, rank, tstart): - global barrier_time - global run - - if not timing: - return 0.0 - if group is not None: - barrier_tstart = time.time() - # dist.barrier(group) - torch.cuda.synchronize(device=device) - barrier_tstop = time.time() - barrier_time[run][rank] += barrier_tstop - barrier_tstart - tstop = 0.0 - if rank == 0: - tstop = time.time() - return tstop - tstart - -def normalize(adj_matrix): - adj_matrix = adj_matrix + torch.eye(adj_matrix.size(0)) - d = torch.sum(adj_matrix, dim=1) - d = torch.rsqrt(d) - d = torch.diag(d) - return torch.mm(d, torch.mm(adj_matrix, d)) - -def block_row(adj_matrix, am_partitions, inputs, weight, rank, size): - n_per_proc = math.ceil(float(adj_matrix.size(1)) / size) - # n_per_proc = int(adj_matrix.size(1) / size) - # am_partitions = list(torch.split(adj_matrix, n_per_proc, dim=1)) - - z_loc = torch.cuda.FloatTensor(n_per_proc, inputs.size(1)).fill_(0) - # z_loc = torch.zeros(adj_matrix.size(0), inputs.size(1)) - - inputs_recv = torch.zeros(inputs.size()) - - part_id = rank % size - - z_loc += torch.mm(am_partitions[part_id].t(), inputs) - - for i in range(1, size): - part_id = (rank + i) % size - - inputs_recv = torch.zeros(am_partitions[part_id].size(0), inputs.size(1)) - - src = (rank + 1) % size - dst = rank - 1 - if dst < 0: - dst = size - 1 - - if rank == 0: - dist.send(tensor=inputs, dst=dst) - dist.recv(tensor=inputs_recv, src=src) - else: - dist.recv(tensor=inputs_recv, src=src) - dist.send(tensor=inputs, dst=dst) - - inputs = inputs_recv.clone() - - # z_loc += torch.mm(am_partitions[part_id], inputs) - z_loc += torch.mm(am_partitions[part_id].t(), inputs) - - - # z_loc = torch.mm(z_loc, weight) - return z_loc - -def outer_product2(inputs, ag, rank, size, group): - global comm_time - global comp_time - global dcomp_time - global op_comm_time - global run - - tstart_comp = start_time(group, rank) - # (H^(l-1))^T * (A * G^l) - grad_weight = torch.mm(inputs, ag) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - tstart_comm = start_time(group, rank) - # reduction on grad_weight low-rank matrices - dist.all_reduce(grad_weight, op=dist.reduce_op.SUM, group=group) - - dur = stop_time(group, rank, tstart_comm) - comm_time[run][rank] += dur - op_comm_time[run][rank] += dur - - return grad_weight - -def broad_func(node_count, am_partitions, inputs, rank, size, row_groups, col_groups, group): - global device - global comm_time - global comp_time - global scomp_time - global bcast_comm_time - global bcast_words - global reduce_comm_time - global run - global replication - - # n_per_proc = math.ceil(float(adj_matrix.size(1)) / size) - n_per_proc = math.ceil(float(node_count) / (size / replication)) - - # z_loc = torch.cuda.FloatTensor(adj_matrix.size(0), inputs.size(1), device=device).fill_(0) - z_loc = torch.cuda.FloatTensor(am_partitions[0].size(0), inputs.size(1), device=device).fill_(0) - # z_loc = torch.zeros(adj_matrix.size(0), inputs.size(1)) - - inputs_recv = torch.cuda.FloatTensor(n_per_proc, inputs.size(1), device=device).fill_(0) - # inputs_recv = torch.zeros(n_per_proc, inputs.size(1)) - - rank_c = rank // replication - rank_col = rank % replication - - stages = size // (replication ** 2) - if rank_col == replication - 1: - stages = (size // replication) - (replication - 1) * stages - - for i in range(stages): - # q = rank_c // (size // (replication ** 2)) * (size // (replication ** 2)) + i - # = q * replication + rank_c // (size // (replication **2)) - q = (rank_col * (size // (replication ** 2)) + i) * replication + rank_col - - q_c = q // replication - - am_partid = rank_col * (size // replication ** 2) + i - - if q == rank: - inputs_recv = inputs.clone() - elif q_c == size // replication - 1: - inputs_recv = torch.cuda.FloatTensor(am_partitions[am_partid].size(1), inputs.size(1), device=device).fill_(0) - # inputs_recv = torch.zeros(list(am_partitions[i].t().size())[1], inputs.size(1)) - tstart_comm = start_time(col_groups[rank_col], rank) - - inputs_recv = inputs_recv.contiguous() - bcast_words[run][rank] += inputs_recv.size(0) * inputs_recv.size(1) - dist.broadcast(inputs_recv, src=q, group=col_groups[rank_col]) - - dur = stop_time(col_groups[rank_col], rank, tstart_comm) - - comm_time[run][rank] += dur - bcast_comm_time[run][rank] += dur - - tstart_comp = start_time(col_groups[rank_col], rank) - - spmm_gpu(am_partitions[am_partid].indices()[0].int(), am_partitions[am_partid].indices()[1].int(), - am_partitions[am_partid].values(), am_partitions[am_partid].size(0), - am_partitions[am_partid].size(1), inputs_recv, z_loc) - - dur = stop_time(col_groups[rank_col], rank, tstart_comp) - comp_time[run][rank] += dur - scomp_time[run][rank] += dur - - z_loc = z_loc.contiguous() - - tstart_comm = start_time(row_groups[rank_c], rank) - dist.all_reduce(z_loc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) - dur = stop_time(row_groups[rank_c], rank, tstart_comm) - - comm_time[run][rank] += dur - reduce_comm_time[run][rank] += dur - - return z_loc - -class GCNFunc(torch.autograd.Function): - @staticmethod - def forward(ctx, inputs, weight, adj_matrix, am_partitions, rank, size, group, row_groups, col_groups, func): - global comm_time - global comp_time - global dcomp_time - global run - - # inputs: H - # adj_matrix: A - # weight: W - # func: sigma - - # adj_matrix = adj_matrix.to_dense() - ctx.save_for_backward(inputs, weight, adj_matrix) - ctx.rank = rank - ctx.size = size - ctx.group = group - ctx.row_groups = row_groups - ctx.col_groups = col_groups - ctx.am_partitions = am_partitions - - ctx.func = func - - # z = block_row(adj_matrix.t(), am_partitions, inputs, weight, rank, size) - z = broad_func(adj_matrix.size(0), am_partitions, inputs, rank, size, row_groups, col_groups, group) - - tstart_comp = start_time(row_groups[0], rank) - - z = torch.mm(z, weight) - - dur = stop_time(row_groups[0], rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - z.requires_grad = True - ctx.z = z - - if activations: - if func is F.log_softmax: - h = func(z, dim=1) - elif func is F.relu: - h = func(z) - else: - h = z - - return h - else: - return z - - @staticmethod - def backward(ctx, grad_output): - global comm_time - global comp_time - global dcomp_time - global run - - inputs, weight, adj_matrix = ctx.saved_tensors - rank = ctx.rank - size = ctx.size - group = ctx.group - row_groups = ctx.row_groups - col_groups = ctx.col_groups - am_partitions = ctx.am_partitions - - func = ctx.func - z = ctx.z - - rank_col = rank % replication - - if activations: - with torch.set_grad_enabled(True): - if func is F.log_softmax: - func_eval = func(z, dim=1) - elif func is F.relu: - func_eval = func(z) - else: - func_eval = z - - sigmap = torch.autograd.grad(outputs=func_eval, inputs=z, grad_outputs=grad_output)[0] - grad_output = sigmap - - # First backprop equation - # ag = outer_product(adj_matrix, grad_output, rank, size, group) - ag = broad_func(adj_matrix.size(0), am_partitions, grad_output, rank, size, row_groups, col_groups, group) - - tstart_comp = start_time(group, rank) - - grad_input = torch.mm(ag, weight.t()) - - dur = stop_time(group, rank, tstart_comp) - comp_time[run][rank] += dur - dcomp_time[run][rank] += dur - - # Second backprop equation (reuses the A * G^l computation) - # grad_weight = outer_product2(inputs.t(), ag, rank, size, group) - grad_weight = outer_product2(inputs.t(), ag, rank, size, col_groups[rank_col]) - - return grad_input, grad_weight, None, None, None, None, None, None, None, None - -def train(inputs, weight1, weight2, adj_matrix, am_partitions, optimizer, data, rank, size, group, row_groups, col_groups): - - outputs = GCNFunc.apply(inputs, weight1, adj_matrix, am_partitions, rank, size, group, row_groups, col_groups, F.relu) - outputs = GCNFunc.apply(outputs, weight2, adj_matrix, am_partitions, rank, size, group, row_groups, col_groups, F.log_softmax) - - optimizer.zero_grad() - - rank_c = rank // replication - node_count = adj_matrix.size(0) - n_per_proc = int(math.ceil(float(node_count) / (size / replication))) - rank_train_mask = torch.split(data.train_mask.bool(), n_per_proc, dim=0)[rank_c] - datay_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] - - # Note: bool type removes warnings, unsure of perf penalty - # loss = F.nll_loss(outputs[data.train_mask.bool()], data.y[data.train_mask.bool()]) - if list(datay_rank[rank_train_mask].size())[0] > 0: - # if datay_rank.size(0) > 0: - loss = F.nll_loss(outputs[rank_train_mask], datay_rank[rank_train_mask]) - # loss = F.nll_loss(outputs, torch.max(datay_rank, 1)[1]) - loss.backward() - else: - fake_loss = (outputs * torch.cuda.FloatTensor(outputs.size(), device=device).fill_(0)).sum() - # fake_loss = (outputs * torch.zeros(outputs.size())).sum() - fake_loss.backward() - - optimizer.step() - - return outputs - -def test(outputs, data, vertex_count, rank): - logits, accs = outputs, [] - for _, mask in data('train_mask', 'val_mask', 'test_mask'): - pred = logits[mask].max(1)[1] - acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - accs.append(acc) - return accs - # logits, accs = outputs, [] - # datay_rank = torch.split(data.y, vertex_count)[rank] - # for _, mask in data('train_mask', 'val_mask', 'test_mask'): - # mask_rank = torch.split(mask, vertex_count)[rank] - # count = mask_rank.nonzero().size(0) - # if count > 0: - # pred = logits[mask_rank].max(1)[1] - # acc = pred.eq(datay_rank[mask_rank]).sum().item() / mask_rank.sum().item() - # # pred = logits[mask].max(1)[1] - # # acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - # else: - # acc = -1 - # accs.append(acc) - # return accs - -def get_proc_groups(rank, size): - global replication - - rank_c = rank // replication - - row_procs = [] - for i in range(0, size, replication): - row_procs.append(list(range(i, i + replication))) - - col_procs = [] - for i in range(replication): - col_procs.append(list(range(i, size, replication))) - - row_groups = [] - for i in range(len(row_procs)): - row_groups.append(dist.new_group(row_procs[i])) - - col_groups = [] - for i in range(len(col_procs)): - col_groups.append(dist.new_group(col_procs[i])) - - return row_groups, col_groups - -# Split a COO into partitions of size n_per_proc -# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. -def split_coo(adj_matrix, node_count, n_per_proc, dim): - vtx_indices = list(range(0, node_count, n_per_proc)) - vtx_indices.append(node_count) - - am_partitions = [] - for i in range(len(vtx_indices) - 1): - am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] - am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] - am_part[dim] -= vtx_indices[i] - am_partitions.append(am_part) - - return am_partitions, vtx_indices - -# Normalize all elements according to KW's normalization rule -def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx): - if not normalization: - return adj_part - - # Scale each edge (u, v) by 1 / (sqrt(u) * sqrt(v)) - # indices = adj_part._indices() - # values = adj_part._values() - - # deg_map = dict() - # for i in range(adj_part._nnz()): - # u = indices[0][i] + row_vtx - # v = indices[1][i] + col_vtx - - # if u.item() in deg_map: - # degu = deg_map[u.item()] - # else: - # degu = (adj_matrix[0] == u).sum().item() - # deg_map[u.item()] = degu - - # if v.item() in deg_map: - # degv = deg_map[v.item()] - # else: - # degv = (adj_matrix[0] == v).sum().item() - # deg_map[v.item()] = degv - - # values[i] = values[i] / (math.sqrt(degu) * math.sqrt(degv)) - - adj_part = adj_part.coalesce() - deg = torch.histc(adj_matrix[0].double(), bins=node_count) - deg = deg.pow(-0.5) - - row_len = adj_part.size(0) - col_len = adj_part.size(1) - - dleft = torch.sparse_coo_tensor([np.arange(0, row_len).tolist(), - np.arange(0, row_len).tolist()], - deg[row_vtx:(row_vtx + row_len)].float(), - size=(row_len, row_len), - requires_grad=False, device=torch.device("cpu")) - - dright = torch.sparse_coo_tensor([np.arange(0, col_len).tolist(), - np.arange(0, col_len).tolist()], - deg[col_vtx:(col_vtx + col_len)].float(), - size=(col_len, col_len), - requires_grad=False, device=torch.device("cpu")) - # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) - ad_ind, ad_val = torch_sparse.spspmm(adj_part._indices(), adj_part._values(), - dright._indices(), dright._values(), - adj_part.size(0), adj_part.size(1), dright.size(1)) - - adj_part_ind, adj_part_val = torch_sparse.spspmm(dleft._indices(), dleft._values(), - ad_ind, ad_val, - dleft.size(0), dleft.size(1), adj_part.size(1)) - - adj_part = torch.sparse_coo_tensor(adj_part_ind, adj_part_val, - size=(adj_part.size(0), adj_part.size(1)), - requires_grad=False, device=torch.device("cpu")) - - return adj_part - -def oned_partition(rank, size, inputs, adj_matrix, data, features, classes, device): - node_count = inputs.size(0) - # n_per_proc = math.ceil(float(node_count) / size) - n_per_proc = math.ceil(float(node_count) / (size / replication)) - - am_partitions = None - am_pbyp = None - - inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) - - rank_c = rank // replication - # Compute the adj_matrix and inputs partitions for this process - # TODO: Maybe I do want grad here. Unsure. - with torch.no_grad(): - # Column partitions - am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) - - proc_node_count = vtx_indices[rank_c + 1] - vtx_indices[rank_c] - am_pbyp, _ = split_coo(am_partitions[rank_c], node_count, n_per_proc, 0) - for i in range(len(am_pbyp)): - if i == size // replication - 1: - last_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(last_node_count, proc_node_count), - requires_grad=False) - - am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank_c]) - else: - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(n_per_proc, proc_node_count), - requires_grad=False) - - am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank_c]) - - for i in range(len(am_partitions)): - proc_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], - torch.ones(am_partitions[i].size(1)), - size=(node_count, proc_node_count), - requires_grad=False) - am_partitions[i] = scale_elements(adj_matrix, am_partitions[i], node_count, 0, vtx_indices[i]) - - input_partitions = torch.split(inputs, math.ceil(float(inputs.size(0)) / (size / replication)), dim=0) - - adj_matrix_loc = am_partitions[rank_c] - inputs_loc = input_partitions[rank_c] - - print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) - print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}", flush=True) - return inputs_loc, adj_matrix_loc, am_pbyp - -def run(rank, size, inputs, adj_matrix, data, features, classes, device): - global epochs - global mid_layer - global timing - global run - - best_val_acc = test_acc = 0 - outputs = None - - # adj_matrix_loc = torch.rand(node_count, n_per_proc) - # inputs_loc = torch.rand(n_per_proc, inputs.size(1)) - - group = dist.new_group(list(range(size))) - row_groups, col_groups = get_proc_groups(rank, size) - - rank_c = rank // replication - rank_col = rank % replication - if rank_c >= (size // replication): - return - - inputs_loc, adj_matrix_loc, am_pbyp = oned_partition(rank, size, inputs, adj_matrix, data, - features, classes, device) - - inputs_loc = inputs_loc.to(device) - adj_matrix_loc = adj_matrix_loc.to(device) - for i in range(len(am_pbyp)): - am_pbyp[i] = am_pbyp[i].t().coalesce().to(device) - - adj_matrix_loc.coalesce() - dist.barrier(group) - - for i in range(run_count): - run = i - torch.manual_seed(0) - weight1_nonleaf = torch.rand(features, mid_layer, requires_grad=True) - weight1_nonleaf = weight1_nonleaf.to(device) - weight1_nonleaf.retain_grad() - - weight2_nonleaf = torch.rand(mid_layer, classes, requires_grad=True) - weight2_nonleaf = weight2_nonleaf.to(device) - weight2_nonleaf.retain_grad() - - weight1 = Parameter(weight1_nonleaf) - weight2 = Parameter(weight2_nonleaf) - - optimizer = torch.optim.Adam([weight1, weight2], lr=0.01) - - total_time[i] = dict() - comm_time[i] = dict() - comp_time[i] = dict() - scomp_time[i] = dict() - dcomp_time[i] = dict() - bcast_comm_time[i] = dict() - bcast_words[i] = dict() - reduce_comm_time[i] = dict() - op_comm_time[i] = dict() - barrier_time[i] = dict() - - total_time[i][rank] = 0.0 - comm_time[i][rank] = 0.0 - comp_time[i][rank] = 0.0 - scomp_time[i][rank] = 0.0 - dcomp_time[i][rank] = 0.0 - bcast_comm_time[i][rank] = 0.0 - bcast_words[i][rank] = 0 - reduce_comm_time[i][rank] = 0.0 - op_comm_time[i][rank] = 0.0 - barrier_time[i][rank] = 0.0 - - # Do not time first epoch - timing_on = timing == True - timing = False - - outputs = train(inputs_loc, weight1, weight2, adj_matrix_loc, am_pbyp, optimizer, data, - rank, size, group, row_groups, col_groups) - if timing_on: - timing = True - - dist.barrier(group) - tstart = time.time() - - # for epoch in range(1, 201): - print(f"Starting training... rank {rank} run {i}", flush=True) - for epoch in range(1, epochs): - outputs = train(inputs_loc, weight1, weight2, adj_matrix_loc, am_pbyp, optimizer, data, - rank, size, group, row_groups, col_groups) - print("Epoch: {:03d}".format(epoch), flush=True) - - # dist.barrier(group) - torch.cuda.synchronize(device=device) - tstop = time.time() - total_time[i][rank] = tstop - tstart - - # Get median runtime according to rank0 and print that run's breakdown - dist.barrier(group) - - if rank == 0: - total_times_r0 = [] - for i in range(run_count): - total_times_r0.append(total_time[i][0]) - - print(f"total_times_r0: {total_times_r0}") - median_run_time = statistics.median(total_times_r0) - median_idx = total_times_r0.index(median_run_time) - median_idx = torch.cuda.LongTensor([median_idx]) - else: - median_idx = torch.cuda.LongTensor([0]) - - dist.barrier(group) - # dist.broadcast(median_idx, src=0, group=group) - median_idx = median_idx.item() - print(f"rank: {rank} median_run: {median_idx}") - print(f"rank: {rank} total_time: {total_time[median_idx][rank]}") - print(f"rank: {rank} comm_time: {comm_time[median_idx][rank]}") - print(f"rank: {rank} comp_time: {comp_time[median_idx][rank]}") - print(f"rank: {rank} scomp_time: {scomp_time[median_idx][rank]}") - print(f"rank: {rank} dcomp_time: {dcomp_time[median_idx][rank]}") - print(f"rank: {rank} bcast_comm_time: {bcast_comm_time[median_idx][rank]}") - print(f"rank: {rank} bcast_words: {bcast_words[median_idx][rank]}") - print(f"rank: {rank} reduce_comm_time: {reduce_comm_time[median_idx][rank]}") - print(f"rank: {rank} op_comm_time: {op_comm_time[median_idx][rank]}") - print(f"rank: {rank} barrier_time: {barrier_time[median_idx][rank]}") - print(f"rank: {rank} {outputs}") - - - if accuracy: - # All-gather outputs to test accuracy - output_parts = [] - # n_per_proc = math.ceil(float(inputs.size(0)) / size) - n_per_proc = math.ceil(float(inputs.size(0)) / (size / replication)) - # print(f"rows: {am_pbyp[-1].size(0)} cols: {classes}", flush=True) - for i in range(size // replication): - output_parts.append(torch.cuda.FloatTensor(n_per_proc, classes, device=device).fill_(0)) - - if outputs.size(0) != n_per_proc: - pad_row = n_per_proc - outputs.size(0) - outputs = torch.cat((outputs, torch.cuda.FloatTensor(pad_row, classes, device=device)), dim=0) - - # dist.all_gather(output_parts, outputs) - dist.all_gather(output_parts, outputs, group=col_groups[rank_col]) - # output_parts[rank] = outputs - output_parts[rank_c] = outputs - - padding = inputs.size(0) - n_per_proc * ((size // replication) - 1) - output_parts[(size // replication) - 1] = output_parts[(size // replication) - 1][:padding,:] - - outputs = torch.cat(output_parts, dim=0) - - train_acc, val_acc, tmp_test_acc = test(outputs, data, am_pbyp[0].size(1), rank) - if val_acc > best_val_acc: - best_val_acc = val_acc - test_acc = tmp_test_acc - log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}' - - print(log.format(900, train_acc, best_val_acc, test_acc)) - return outputs - -def rank_to_devid(rank, acc_per_rank): - return rank % acc_per_rank - -def init_process(rank, size, inputs, adj_matrix, data, features, classes, device, outputs, fn): - run_outputs = fn(rank, size, inputs, adj_matrix, data, features, classes, device) - if outputs is not None: - outputs[rank] = run_outputs.detach() - -def main(): - global device - global graphname - - print(socket.gethostname()) - seed = 0 - - if not download: - mp.set_start_method('spawn', force=True) - outputs = None - if "OMPI_COMM_WORLD_RANK" in os.environ.keys(): - os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] - - # Initialize distributed environment with SLURM - if "SLURM_PROCID" in os.environ.keys(): - os.environ["RANK"] = os.environ["SLURM_PROCID"] - - if "SLURM_NTASKS" in os.environ.keys(): - os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] - - if "MASTER_ADDR" not in os.environ.keys(): - os.environ["MASTER_ADDR"] = "127.0.0.1" - - os.environ["MASTER_PORT"] = "1234" - dist.init_process_group(backend='nccl') - rank = dist.get_rank() - size = dist.get_world_size() - print("Processes: " + str(size)) - - # device = torch.device('cpu') - devid = rank_to_devid(rank, acc_per_rank) - device = torch.device('cuda:{}'.format(devid)) - torch.cuda.set_device(device) - curr_devid = torch.cuda.current_device() - # print(f"curr_devid: {curr_devid}", flush=True) - devcount = torch.cuda.device_count() - - if graphname == "Cora": - # path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data') - dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures()) - data = dataset[0] - data = data.to(device) - data.x.requires_grad = True - inputs = data.x.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - edge_index = data.edge_index - num_features = dataset.num_features - num_classes = dataset.num_classes - elif graphname == "Reddit": - path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - dataset = Reddit(path, T.NormalizeFeatures()) - data = dataset[0] - data = data.to(device) - data.x.requires_grad = True - inputs = data.x.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - edge_index = data.edge_index - num_features = dataset.num_features - num_classes = dataset.num_classes - elif graphname == 'Amazon': - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/Amazon/processed/data.pt") - edge_index = edge_index.t_() - print(f"Done loading coo", flush=True) - # edge_index = edge_index.t_() - # n = 9430088 - n = 14249639 - # n = 14249640 - num_features = 300 - num_classes = 24 - # mid_layer = 24 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1).long() - data.train_mask = torch.ones(n).long() - # edge_index = edge_index.to(device) - print(f"edge_index.size: {edge_index.size()}", flush=True) - print(f"edge_index: {edge_index}", flush=True) - data = data.to(device) - # inputs = inputs.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - elif graphname == 'subgraph3': - print(f"Loading coo...", flush=True) - # edge_index = torch.load("../data/subgraph3/processed/data.pt") - edge_index = torch.load("../data/protein/processed/protein.pt") - edge_index = edge_index.t_() - print(f"Done loading coo", flush=True) - n = 8745542 - num_features = 128 - # mid_layer = 512 - # mid_layer = 64 - num_classes = 256 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1).long() - data.train_mask = torch.ones(n).long() - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - - if download: - exit() - - if normalization: - adj_matrix, _ = add_remaining_self_loops(edge_index, num_nodes=inputs.size(0)) - else: - adj_matrix = edge_index - - - init_process(rank, size, inputs, adj_matrix, data, num_features, num_classes, device, outputs, - run) - - if outputs is not None: - return outputs[0] - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument("--local_rank", type=int) - parser.add_argument("--accperrank", type=int) - parser.add_argument("--epochs", type=int) - parser.add_argument("--graphname", type=str) - parser.add_argument("--timing", type=str) - parser.add_argument("--midlayer", type=int) - parser.add_argument("--runcount", type=int) - parser.add_argument("--replication", type=int) - parser.add_argument("--normalization", type=str) - parser.add_argument("--activations", type=str) - parser.add_argument("--accuracy", type=str) - parser.add_argument("--download", type=bool) - - args = parser.parse_args() - print(args) - - acc_per_rank = args.accperrank - - epochs = args.epochs - graphname = args.graphname - timing = args.timing == "True" - mid_layer = args.midlayer - run_count = args.runcount - normalization = args.normalization == "True" - activations = args.activations == "True" - accuracy = args.accuracy == "True" - replication = args.replication - download = args.download - - if not download: - if (epochs is None) or (graphname is None) or (timing is None) or (mid_layer is None) or (run_count is None): - print(f"Error: missing argument {epochs} {graphname} {timing} {mid_layer} {run_count}") - exit() - - print(f"Arguments: epochs: {epochs} graph: {graphname} timing: {timing} mid: {mid_layer} norm: {normalization} act: {activations} acc: {accuracy} runs: {run_count} rep: {replication}") - - print(main()) diff --git a/gcn_distr_2d.py b/gcn_distr_2d.py deleted file mode 100644 index 1906b87..0000000 --- a/gcn_distr_2d.py +++ /dev/null @@ -1,1575 +0,0 @@ -import os -import os.path as osp -import argparse - -import math - -import torch -import torch_sparse -import torch.distributed as dist - -from torch_geometric.data import Data, Dataset -from torch_geometric.datasets import Planetoid, PPI -from reddit import Reddit -from torch_geometric.nn import GCNConv, ChebConv # noqa -from torch_geometric.utils import ( - add_remaining_self_loops, - to_dense_adj, - dense_to_sparse, - to_scipy_sparse_matrix -) -import torch_geometric.transforms as T - -import torch.multiprocessing as mp - -from torch.multiprocessing import Manager, Process - -import statistics - -from torch.nn import Parameter -import torch.nn.functional as F - -from torch_scatter import scatter_add - -import socket -import time -import numpy as np - -from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu - -# comp_time = 0.0 -# comm_time = 0.0 -# summa_sparse_bcast1 = 0.0 -# summa_sparse_bcast1_words = 0 -# summa_sparse_bcast2_words = 0 -# summa_sparse_bcast2 = 0.0 -# summa_sparse_bcast2_fwd = 0.0 -# summa_sparse_bcast2_bwd = 0.0 -# summa_bcast1 = 0.0 -# summa_bcast2 = 0.0 -# summa_sparse_comp = 0.0 -# summa_comp = 0.0 -# summa_loc_bcast = 0.0 -# fwd_time = 0.0 -# bwd_time = 0.0 -# transpose_time = 0.0 -# grad_weight_time = 0.0 -# loss_calc_time = 0.0 -# summa_sparse_time = 0.0 -# summa_time = 0.0 -# summa_loc_time = 0.0 -total_time = dict() -comp_time = dict() -comm_time = dict() -summa_sparse_bcast1 = dict() -summa_sparse_bcast1_words = dict() -summa_sparse_bcast2_words = dict() -summa_sparse_bcast2 = dict() -summa_sparse_bcast2_fwd = dict() -summa_sparse_bcast2_bwd = dict() -summa_bcast1 = dict() -summa_bcast2 = dict() -summa_sparse_comp = dict() -summa_comp = dict() -summa_loc_bcast = dict() -fwd_time = dict() -bwd_time = dict() -transpose_time = dict() -grad_weight_time = dict() -loss_calc_time = dict() -summa_sparse_time = dict() -summa_time = dict() -summa_loc_time = dict() - -epochs = 0 -graphname = "" -mid_layer = 0 -timing = False -normalization = False -activations = False -accuracy = False -no_occur_val = 42.1234 -run_count = 0 -run = 0 -download = False - -def sync_and_sleep(rank, device): - torch.cuda.synchronize(device=device) - print(f"Sleeping rank {rank}", flush=True) - time.sleep(20) - print(f"Done sleeping rank {rank}", flush=True) - -def normalize(adj_matrix): - adj_matrix = adj_matrix + torch.eye(adj_matrix.size(0)) - d = torch.sum(adj_matrix, dim=1) - d = torch.rsqrt(d) - d = torch.diag(d) - return torch.mm(d, torch.mm(adj_matrix, d)) - -def start_time(group, rank): - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - if not timing: - return 0.0 - if group is not None: - # dist.barrier(group) - torch.cuda.synchronize(device=device) - tstart = 0.0 - if rank == 0: - tstart = time.time() - return tstart - -def stop_time(group, rank, tstart): - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - if not timing: - return 0.0 - if group is not None: - # dist.barrier(group) - torch.cuda.synchronize(device=device) - tstop = 0.0 - if rank == 0: - tstop = time.time() - return tstop - tstart - -def transpose(mat, row, col, height, width, size, acc_per_rank, transpose_group): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - rank = row * proc_col + col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - rank_t = col * proc_row + row - - if rank == rank_t: - return mat.t() - - # height_recv = math.ceil(float(width) / proc_row) - # width_recv = math.ceil(float(height) / proc_col) - height_recv = width // proc_row - width_recv = height // proc_col - - if row == proc_row - 1: - # height_recv -= proc_row * height_recv - width - height_recv = width - height_recv * (proc_row - 1) - - if col == proc_col - 1: - # width_recv -= proc_col * width_recv - height - width_recv = height - width_recv * (proc_col - 1) - - mat_recv = torch.cuda.FloatTensor(height_recv, width_recv, device=device) - - # if rank < rank_t: - # dist.send(tensor=mat.t().contiguous(), dst=rank_t) - # dist.recv(tensor=mat_recv, src=rank_t) - # else: - # dist.recv(tensor=mat_recv, src=rank_t) - # dist.send(tensor=mat.t().contiguous(), dst=rank_t) - - # transpose_group = dist.new_group([rank, rank_t]) - - mat_recvs = [mat.t().contiguous(), mat_recv] - - if rank < rank_t: - dist.broadcast(mat_recvs[0], src=rank, group=transpose_group) - dist.broadcast(mat_recvs[1], src=rank_t, group=transpose_group) - # dist.broadcast_multigpu([mat_recvs[0]], src=rank, group=transpose_group) - # dist.broadcast_multigpu([mat_recvs[1]], src=rank_t, group=transpose_group) - else: - dist.broadcast(mat_recvs[1], src=rank_t, group=transpose_group) - dist.broadcast(mat_recvs[0], src=rank, group=transpose_group) - # dist.broadcast_multigpu([mat_recvs[1]], src=rank_t, group=transpose_group) - # dist.broadcast_multigpu([mat_recvs[0]], src=rank, group=transpose_group) - - return mat_recvs[1] - -def summa(adj_matrix, inputs, rank, row, col, size, acc_per_rank, row_groups, col_groups, height, - middim, width): - - global comm_time - global comp_time - - global summa_bcast1 - global summa_bcast2 - - global summa_comp - global summa_time - global run - - # tstart_summa_time = start_time(row_groups[0], rank) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - # height_per_proc = math.ceil(float(height) / proc_row) - # width_per_proc = math.ceil(float(width) / proc_col) - # # TODO: Not sure how to handle this w/o square grid - # middim_per_proc = math.ceil(float(middim) / proc_row) - height_per_proc = height // proc_row - width_per_proc = width // proc_col - # TODO: Not sure how to handle this w/o square grid - middim_per_proc = middim // proc_row - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - if row == proc_row - 1: - # height_per_proc -= proc_row * height_per_proc - height - height_per_proc = height - height_per_proc * (proc_row - 1) - - if col == proc_col - 1: - # width_per_proc -= proc_col * width_per_proc - width - width_per_proc = width - width_per_proc * (proc_col - 1) - - acol_tens = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=device) - brow_tens = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=device) - - acol = acol_tens - brow = brow_tens - - z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - - for k in range(proc_col): - - row_src_rank = k + proc_col * row - col_src_rank = k * proc_col + col - - if k == proc_col - 1: - # middim_per_proc -= proc_col * middim_per_proc - middim - middim_per_proc = middim - middim_per_proc * (proc_col - 1) - # acol_tens = acol_tens[:,:middim_per_proc] - # brow_tens = brow_tens[:middim_per_proc] - acol_tens = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=device) - brow_tens = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=device) - - if row_src_rank == rank: - acol = adj_matrix - else: - acol = acol_tens - # acol = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=device) - - tstart = start_time(row_groups[row], rank) - - acol = acol.contiguous() - # dist.broadcast_multigpu([acol], row_src_rank, row_groups[row]) - dist.broadcast(acol, row_src_rank, row_groups[row]) - - dur = stop_time(row_groups[row], rank, tstart) - comm_time[run][rank] += dur - summa_bcast1[run][rank] += dur - - if col_src_rank == rank: - brow = inputs - else: - brow = brow_tens - # brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=device) - - tstart = start_time(col_groups[col], rank) - - brow = brow.contiguous() - # dist.broadcast_multigpu([brow], col_src_rank, col_groups[col]) - dist.broadcast(brow, col_src_rank, col_groups[col]) - - dur = stop_time(col_groups[col], rank, tstart) - comm_time[run][rank] += dur - summa_bcast2[run][rank] += dur - - # tstart = start_time(row_groups[0], rank) - tstart = start_time(None, rank) - - z_loc += torch.mm(acol.float(), brow) - - # dur = stop_time(row_groups[0], rank, tstart) - dur = stop_time(None, rank, tstart) - comp_time[run][rank] += dur - summa_comp[run][rank] += dur - - # summa_time += stop_time(row_groups[0], rank, tstart_summa_time) - return z_loc - -def summa_sparse(adj_matrix, inputs, rank, row, col, size, acc_per_rank, row_groups, col_groups, - height, middim, width): - - global comm_time - global comp_time - - global summa_sparse_bcast1 - global summa_sparse_bcast2 - - global summa_sparse_bcast1_words - global summa_sparse_bcast2_words - - global summa_sparse_comp - global summa_sparse_time - global run - - # tstart_summa_sparse_time = start_time(row_groups[0], rank) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - # height_per_proc = math.ceil(float(height) / proc_row) - # width_per_proc = math.ceil(float(width) / proc_col) - - # # TODO: Not sure how to handle this w/o square grid - # middim_per_proc = math.ceil(float(middim) / proc_col) - height_per_proc = height // proc_row - width_per_proc = width // proc_col - - # TODO: Not sure how to handle this w/o square grid - middim_per_proc = middim // proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - if row == proc_row - 1: - # height_per_proc -= proc_row * height_per_proc - height - height_per_proc = height - height_per_proc * (proc_row - 1) - - if col == proc_col - 1: - # width_per_proc -= proc_col * width_per_proc - width - width_per_proc = width - width_per_proc * (proc_col - 1) - - # acol = torch.cuda.sparse.FloatTensor(height_per_proc, middim_per_proc, device=device) - - # brow = torch.FloatTensor(middim_per_proc, width_per_proc) - - z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - - for k in range(proc_col): - - row_src_rank = k + proc_col * row - col_src_rank = k * proc_col + col - - if k == proc_col - 1: - # middim_per_proc -= proc_col * middim_per_proc - middim - middim_per_proc = middim - middim_per_proc * (proc_col - 1) - - if row_src_rank == rank: - # acol = adj_matrix.clone() - acol_indices_len = torch.cuda.LongTensor( - [adj_matrix.indices().contiguous()[0].size(0)], - device=device) - acol_values_len = torch.cuda.LongTensor([adj_matrix.values().contiguous().size(0)], - device=device) - else: - # acol = torch.sparse.FloatTensor(height_per_proc, middim_per_proc) - acol_indices_len = torch.cuda.LongTensor([0], device=device) - acol_values_len = torch.cuda.LongTensor([0], device=device) - - dist.broadcast(acol_indices_len, row_src_rank, row_groups[row]) - # dist.broadcast_multigpu([acol_indices_len], row_src_rank, row_groups[row]) - - acol_indices_len = acol_indices_len.item() # nnz - # acol_values_len = acol_values_len.item() - acol_values_len = acol_indices_len - - if row_src_rank == rank: - acol_indices = adj_matrix.indices().contiguous().long() - acol_values = adj_matrix.values().contiguous().float() - else: - acol_indices = torch.cuda.LongTensor(2, acol_indices_len, device=device).fill_(0) - acol_values = torch.cuda.FloatTensor(acol_values_len, device=device).fill_(0) - - - acol = torch.cat((acol_indices.float(), acol_values.unsqueeze(0)), dim=0).contiguous() - - tstart = start_time(row_groups[row], rank) - - # dist.broadcast_multigpu([acol], row_src_rank, row_groups[row]) - dist.broadcast(acol, row_src_rank, row_groups[row]) - - dur = stop_time(row_groups[row], rank, tstart) - comm_time[run][rank] += dur - summa_sparse_bcast1[run][rank] += dur - if rank == 0: - summa_sparse_bcast1_words[run][rank] += 3 * acol_values_len - - acol_indices = acol[:2].long() - acol_values = acol[2].squeeze(0) - - if row_src_rank == rank: - acol = adj_matrix - else: - acol = sparse_coo_tensor_gpu(acol_indices, acol_values, - torch.Size([height_per_proc, middim_per_proc])) - - - if col_src_rank == rank: - brow = inputs - else: - brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=device) - - brow = brow.contiguous() - - tstart = start_time(row_groups[0], rank) - # tstart = start_time(col_groups[col], rank) - - # dist.broadcast_multigpu([brow], col_src_rank, col_groups[col]) - dist.broadcast(brow, col_src_rank, col_groups[col]) - - dur = stop_time(row_groups[0], rank, tstart) - # dur = stop_time(col_groups[col], rank, tstart) - - comm_time[run][rank] += dur - summa_sparse_bcast2[run][rank] += dur - if rank == 0: - summa_sparse_bcast2_words[run][rank] += brow.size(0) * brow.size(1) - - # tstart = start_time(row_groups[0], rank) - tstart = start_time(None, rank) - - spmm_gpu(acol_indices[0].int(), acol_indices[1].int(), acol_values, - height_per_proc, middim_per_proc, brow, z_loc) - - # dur = stop_time(row_groups[0], rank, tstart) - dur = stop_time(None, rank, tstart) - # dur = stop_time(col_groups[col], rank, tstart) - comp_time[run][rank] += dur - summa_sparse_comp[run][rank] += dur - - # summa_sparse_time += stop_time(row_groups[0], rank, tstart_summa_sparse_time) - return z_loc - -def summa_loc(mata, matb, rank, row, col, size, acc_per_rank, row_groups, col_groups, - height, middim, width): - - global comm_time - global comp_time - - global summa_loc_bcast - global summa_loc_time - global run - - # tstart_summa_loc_time = start_time(row_groups[0], rank) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - # height_per_proc = math.ceil(float(height) / proc_row) - # width_per_proc = math.ceil(float(width) / proc_col) - # # TODO: Not sure how to handle this w/o square grid - # middim_per_proc = math.ceil(float(middim) / proc_row) - height_per_proc = height // proc_row - width_per_proc = width // proc_col - # TODO: Not sure how to handle this w/o square grid - middim_per_proc = middim // proc_row - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - if row == proc_row - 1: - # height_per_proc -= proc_row * height_per_proc - height - height_per_proc = height - height_per_proc * (proc_row - 1) - - # if col == proc_col - 1: - # width_per_proc -= proc_col * width_per_proc - width - - width_per_proc = matb[rank].size(1) - - acol_tens = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=device) - brow_tens = torch.FloatTensor(middim_per_proc, width_per_proc) - - acol = acol_tens - brow = brow_tens - - z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - - for k in range(proc_col): - - row_src_rank = k + proc_col * row - col_src_rank = k * proc_col + col - - if k == proc_col - 1: - middim_per_proc -= proc_col * middim_per_proc - middim - - if row_src_rank == rank: - acol = mata - else: - acol = acol_tens - acol = torch.cuda.FloatTensor(height_per_proc, matb[col_src_rank].size(0), - device=device) - - tstart = start_time(row_groups[row], rank) - - acol = acol.contiguous() - # dist.broadcast_multigpu([acol], row_src_rank, row_groups[row]) - dist.broadcast(acol, row_src_rank, row_groups[row]) - - dur = stop_time(row_groups[row], rank, tstart) - comm_time[run][rank] += dur - summa_loc_bcast[run][rank] += dur - - # if col_src_rank == rank: - # brow = matb.clone() - # else: - # brow = torch.FloatTensor(middim_per_proc, width_per_proc) - - # dist.broadcast(brow, col_src_rank, col_groups[col]) - - brow = matb[col_src_rank] - - # tstart = start_time(row_groups[0], rank) - tstart = start_time(None, rank) - - z_loc += torch.mm(acol, brow) - - # dur = stop_time(row_groups[0], rank, tstart) - dur = stop_time(None, rank, tstart) - comp_time[run][rank] += dur - - # summa_loc_time += stop_time(row_groups[0], rank, tstart_summa_loc_time) - return z_loc - -def get_proc_groups(rank, size, group): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - - row_groups = [] - col_groups = [] - - for i in range(proc_row): - # dist.barrier(group) - row_groups.append(dist.new_group(list(range(i * proc_col, i * proc_col + proc_col)))) - - # dist.barrier(group) - for i in range(proc_col): - # dist.barrier(group) - col_groups.append(dist.new_group(list(range(i, size, proc_row)))) - - return row_groups, col_groups - -def dist_log_softmax(z, rank, size, acc_per_rank, group): - torch.set_printoptions(edgeitems=4) - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - maxes = torch.max(z, dim=1, keepdim=True)[0] - maxes_recv = [] - for i in range(proc_col): - maxes_recv.append(torch.cuda.FloatTensor(maxes.size(), device=device)) - - # dist.all_reduce(maxes, op=dist.reduce_op.MAX, group=group) - dist.all_gather(maxes_recv, maxes, group=group) - maxes_recv[rank_col] = maxes - maxes = torch.max(torch.cat(maxes_recv, dim=1), dim=1, keepdim=True)[0] - - h = torch.exp(z - maxes) - sm_sum = torch.sum(h, dim=1, keepdim=True) - - sm_sum_recv = [] - for i in range(proc_col): - sm_sum_recv.append(torch.cuda.FloatTensor(sm_sum.size(), device=device)) - - # dist.all_reduce(sm_sum, op=dist.reduce_op.SUM, group=group) - dist.all_gather(sm_sum_recv, sm_sum, group=group) - sm_sum_recv[rank_col] = sm_sum - sm_sum = torch.sum(torch.cat(sm_sum_recv, dim=1), dim=1, keepdim=True) - sm_sum = torch.log(sm_sum) - h = z - maxes - sm_sum - return h - -def dist_log_softmax2(z, rank, size, width, acc_per_rank, group, grad_output): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - chunk_sizes_col = [] - width_per_col = width // proc_col - - for i in range(proc_col): - if i == proc_col - 1: - chunk_sizes_col.append(width - width_per_col * (proc_col - 1)) - else: - chunk_sizes_col.append(width_per_col) - - width_per_proc = width - width_per_col * (proc_col - 1) - if z.size(1) != width_per_proc: - z = torch.cat((z, torch.cuda.FloatTensor(z.size(0), width_per_proc - z.size(1))), dim=1) - - z_recv = [] - for i in range(proc_col): - z_recv.append(torch.cuda.FloatTensor(z.size())) - - dist.all_gather(z_recv, z, group=group) - z_recv[rank_col] = z - - for i in range(proc_col - 1): - pad_col = width // proc_col - z_recv[i] = z_recv[i][:,:pad_col] - - z = torch.cat(z_recv, dim=1) - - if grad_output is not None: - if grad_output.size(1) != width_per_proc: - grad_output = torch.cat((grad_output, - torch.cuda.FloatTensor(grad_output.size(0), - width_per_proc - grad_output.size(1))), - dim=1) - - grad_output_recv = [] - for i in range(proc_col): - grad_output_recv.append(torch.cuda.FloatTensor(grad_output.size())) - - dist.all_gather(grad_output_recv, grad_output, group=group) - grad_output_recv[rank_col] = grad_output - - for i in range(proc_col - 1): - pad_col = width // proc_col - grad_output_recv[i] = grad_output_recv[i][:,:pad_col] - - grad_output = torch.cat(grad_output_recv, dim=1) - - maxes = torch.max(z, dim=1, keepdim=True)[0] - h = torch.exp(z - maxes) - sm_sum = torch.sum(h, dim=1, keepdim=True) - sm_sum = torch.log(sm_sum) - - h = z - maxes - sm_sum - - # if h.requires_grad: - # if rank_col == 0: - # sm_sigma = torch.autograd.grad(outputs=h, inputs=z, - # grad_outputs=grad_output)[0] - # print(f"rank: {rank} sm_sigma: {sm_sigma}", flush=True) - # else: - # sm_sigma = torch.autograd.grad(outputs=h, inputs=z, - # grad_outputs=grad_output)[0] - # print(f"rank: {rank} sm_sigma: {sm_sigma}", flush=True) - - # Only works for P = 4 - # if rank_col == 0: - # return h[:,:width_per_proc] - # else: - # return h[:,width_per_proc:] - return h, z, grad_output - -class GCNFunc(torch.autograd.Function): - @staticmethod - def forward(ctx, inputs, weight, node_count, adj_matrix, am_partitions, rank, size, - acc_per_rank, group, row_groups, col_groups, transpose_group, func): - # inputs: H - # adj_matrix: A - # weight: W - # func: sigma - - global summa_sparse_bcast2 - global summa_sparse_bcast2_fwd - global fwd_time - global grad_weight_time - global run - - # tstart = start_time(row_groups[0], rank) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - ctx.save_for_backward(inputs, weight, adj_matrix) - ctx.node_count = node_count - ctx.rank = rank - ctx.size = size - ctx.acc_per_rank = acc_per_rank - ctx.group = group - ctx.row_groups = row_groups - ctx.col_groups = col_groups - ctx.transpose_group = transpose_group - - ctx.func = func - - adj_matrix_t = adj_matrix # Only true for undirected graphs - - tmp_summa_sparse_bcast2 = summa_sparse_bcast2[run][rank] - - # TODO: will need to change height argument when n % sqrt(P) != 0 and non-square grid - z = summa_sparse(adj_matrix_t, inputs, rank, rank_row, rank_col, size, acc_per_rank, - row_groups, col_groups, node_count, node_count, weight.size(0)) - - # tstart_grad_weight = start_time(row_groups[0], rank) - chunk_sizes_row = [] - chunk_sizes_col = [] - weight_per_row = weight.size(0) // proc_row - weight_per_col = weight.size(1) // proc_col - for i in range(proc_row): - if i == proc_row - 1: - chunk_sizes_row.append(weight.size(0) - weight_per_row * (proc_row - 1)) - else: - chunk_sizes_row.append(weight_per_row) - - for i in range(proc_col): - if i == proc_col - 1: - chunk_sizes_col.append(weight.size(1) - weight_per_col * (proc_col - 1)) - else: - chunk_sizes_col.append(weight_per_col) - - # weight_rows = torch.split(weight, math.ceil(float(weight.size(0)) / proc_row), dim=0) - weight_rows = torch.split(weight, chunk_sizes_row, dim=0) - weight_parts = [] - for i in weight_rows: - # weight_cols = torch.split(i, math.ceil(float(weight.size(1)) / proc_col), dim=1) - weight_cols = torch.split(i, chunk_sizes_col, dim=1) - weight_parts.extend(weight_cols) - # grad_weight_time += stop_time(row_groups[0], rank, tstart_grad_weight) - - # z = torch.mm(z, weight) - z = summa_loc(z, weight_parts, rank, rank_row, rank_col, size, acc_per_rank, row_groups, - col_groups, node_count, weight.size(0), weight.size(1)) - - z.requires_grad = True - ctx.z = z - - summa_sparse_bcast2_fwd[run][rank] += summa_sparse_bcast2[run][rank] - tmp_summa_sparse_bcast2 - - if activations: - if func is F.log_softmax: - h = dist_log_softmax(z, rank, size, acc_per_rank, row_groups[rank_row]) - elif func is F.relu: - h = func(z) - else: - h = z - return h - else: - return z - - # dur = stop_time(row_groups[0], rank, tstart) - # fwd_time += dur - - # return z - - @staticmethod - def backward(ctx, grad_output): - global summa_sparse_bcast2 - global summa_sparse_bcast2_bwd - global bwd_time - global transpose_time - global grad_weight_time - global run - - inputs, weight, adj_matrix = ctx.saved_tensors - rank = ctx.rank - size = ctx.size - acc_per_rank = ctx.acc_per_rank - group = ctx.group - row_groups = ctx.row_groups - col_groups = ctx.col_groups - transpose_group = ctx.transpose_group - node_count = ctx.node_count - - func = ctx.func - z = ctx.z - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - # tstart = start_time(row_groups[0], rank) - - if activations: - with torch.set_grad_enabled(True): - if func is F.log_softmax: - # func_eval = dist_log_softmax2(z, rank, size, acc_per_rank, row_groups[rank_row]) - func_eval, z_gathered, go_gathered = dist_log_softmax2(z, rank, size, - weight.size(1), - acc_per_rank, - row_groups[rank_row], grad_output) - width = z_gathered.size(1) - - sigmap = torch.autograd.grad(outputs=func_eval, inputs=z_gathered, - grad_outputs=go_gathered)[0] - - chunk_sizes_col = [] - sigmap_per_col = width // proc_col - - for i in range(proc_col): - if i == proc_col - 1: - chunk_sizes_col.append(width - sigmap_per_col * (proc_col - 1)) - else: - chunk_sizes_col.append(sigmap_per_col) - - grad_output = sigmap.split(chunk_sizes_col, dim=1)[rank_col] - del z_gathered - del go_gathered - elif func is F.relu: - func_eval = func(z) - sigmap = torch.autograd.grad(outputs=func_eval, inputs=z,grad_outputs=grad_output)[0] - grad_output = sigmap - else: - func_eval = z - sigmap = torch.autograd.grad(outputs=func_eval, inputs=z,grad_outputs=grad_output)[0] - grad_output = sigmap - - - tmp_summa_sparse_bcast2 = summa_sparse_bcast2[run][rank] - - # First backprop equation - # TODO: will need to change height argument when n % sqrt(P) != 0 and non-square grid - ag = summa_sparse(adj_matrix, grad_output, rank, rank_row, rank_col, size, acc_per_rank, - row_groups, col_groups, node_count, node_count, weight.t().size(0)) - - # tstart_grad_weight = start_time(row_groups[0], rank) - chunk_sizes_row = [] - chunk_sizes_col = [] - weight_per_row = weight.t().size(0) // proc_row - weight_per_col = weight.t().size(1) // proc_col - for i in range(proc_row): - if i == proc_row - 1: - chunk_sizes_row.append(weight.t().size(0) - weight_per_row * (proc_row - 1)) - else: - chunk_sizes_row.append(weight_per_row) - - for i in range(proc_col): - if i == proc_col - 1: - chunk_sizes_col.append(weight.t().size(1) - weight_per_col * (proc_col - 1)) - else: - chunk_sizes_col.append(weight_per_col) - # weight_rows = torch.split(weight.t(), math.ceil(float(weight.t().size(0)) / proc_row), - weight_rows = torch.split(weight.t(), chunk_sizes_row, dim=0) - - weight_parts = [] - for i in weight_rows: - # weight_cols = torch.split(i, math.ceil(float(weight.t().size(1)) / proc_col), dim=1) - weight_cols = torch.split(i, chunk_sizes_col, dim=1) - weight_parts.extend(weight_cols) - - # grad_input = torch.mm(ag, weight.t()) - grad_input = summa_loc(ag, weight_parts, rank, rank_row, rank_col, size, acc_per_rank, - row_groups, col_groups, node_count, weight.t().size(0), - weight.t().size(1)) - # grad_weight_time += stop_time(row_groups[0], rank, tstart_grad_weight) - - # Second backprop equation (reuses the A * G^l computation) - # col_groups twice because of transpose - # TODO: will need to change height argument when n % sqrt(P) != 0 and non-square grid - - # tstart_transpose = start_time(row_groups[0], rank) - tstart_transpose = start_time(transpose_group, rank) - inputs_t = transpose(inputs, rank_row, rank_col, node_count, weight.size(0), size, - acc_per_rank, transpose_group) - # transpose_time[run][rank] += stop_time(row_groups[0], rank, tstart_transpose) - transpose_time[run][rank] += stop_time(transpose_group, rank, tstart_transpose) - - grad_weight = summa(inputs_t, ag, rank, rank_row, rank_col, size, acc_per_rank, row_groups, - col_groups, weight.size(0), node_count, weight.size(1)) - - # tstart_grad_weight = start_time(row_groups[0], rank) - # Collect grad_weight's across processes - grad_weight_recv = [] - max_row_chunk = max(chunk_sizes_col) #transpose - max_col_chunk = max(chunk_sizes_row) - for i in range(size): - grad_weight_recv.append(torch.cuda.FloatTensor( - max_row_chunk, - max_col_chunk, - device=device)) - - # pad_row = math.ceil(float(weight.size(0)) / proc_row) - grad_weight.size(0) - # pad_col = math.ceil(float(weight.size(1)) / proc_col) - grad_weight.size(1) - pad_row = max_row_chunk - grad_weight.size(0) - pad_col = max_col_chunk - grad_weight.size(1) - - # TODO: make this part less hacky - grad_weight = torch.cat((grad_weight, - torch.cuda.FloatTensor(pad_row, grad_weight.size(1), device=device).fill_(no_occur_val)), - dim=0) - grad_weight = torch.cat((grad_weight, - torch.cuda.FloatTensor(grad_weight.size(0), pad_col, device=device).fill_(no_occur_val)), - dim=1) - - dist.all_gather(grad_weight_recv, grad_weight) - # dist.all_gather_multigpu([grad_weight_recv], [grad_weight]) - - # for i in range(size): - # if rank == i: - # grad_weight_recv[i] = grad_weight - # dist.broadcast(grad_weight_recv[i], i, group) - # grad_weight_recv[0] = grad_weight - - for i in range(len(grad_weight_recv)): - grad_weight_recv[i] = grad_weight_recv[i][(grad_weight_recv[i][:, 0] != no_occur_val) - .nonzero().squeeze(1)] - - grad_weight_recv_t = grad_weight_recv[i].t() - grad_weight_recv_t = grad_weight_recv_t[(grad_weight_recv_t[:, 0] != no_occur_val) - .nonzero().squeeze(1)] - - grad_weight_recv[i] = grad_weight_recv_t.t() - - grad_weight_fin = torch.cuda.FloatTensor(device=device) - for i in range(proc_row): - grad_weight_row = torch.cuda.FloatTensor(device=device) - for j in range(proc_col): - rank_wt = i * proc_row + j - grad_weight_row = torch.cat((grad_weight_row, grad_weight_recv[rank_wt]), dim=1) - grad_weight_fin = torch.cat((grad_weight_fin, grad_weight_row), dim=0) - - summa_sparse_bcast2_bwd[run][rank] += summa_sparse_bcast2[run][rank] - tmp_summa_sparse_bcast2 - - # dur = stop_time(row_groups[0], rank, tstart) - # bwd_time += dur - - # grad_weight_time += stop_time(row_groups[0], rank, tstart_grad_weight) - - return grad_input, grad_weight_fin, None, None, None, None, None, None, None, None, None, None, None - -def train(inputs, weight1, weight2, node_count, adj_matrix, am_partitions, optimizer, data, rank, - size, acc_per_rank, group, row_groups, col_groups, transpose_group): - - global loss_calc_time - global run - - outputs = GCNFunc.apply(inputs, weight1, node_count, adj_matrix, am_partitions, rank, size, - acc_per_rank, group, row_groups, col_groups, transpose_group, - F.relu) - - outputs = GCNFunc.apply(outputs, weight2, node_count, adj_matrix, am_partitions, rank, size, - acc_per_rank, group, row_groups, col_groups, transpose_group, - F.log_softmax) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - optimizer.zero_grad() - rank_train_mask = torch.split(data.train_mask, outputs.size(0), dim=0)[rank_row] - datay_rank = torch.split(data.y, outputs.size(0), dim=0)[rank_row] - - total_classes = weight2.size(1) - # class_per_rank = math.ceil(float(total_classes) / proc_col) - class_per_rank = total_classes // proc_col - - min_class = rank_col * class_per_rank - max_class = min((rank_col + 1) * class_per_rank, total_classes) - if rank_col == proc_col - 1: - max_classes = total_classes - - - # Note: bool type removes warnings, unsure of perf penalty - # loss = F.nll_loss(outputs[data.train_mask.bool()], data.y[data.train_mask.bool()]) - if list(datay_rank[rank_train_mask].size())[0] > 0: - # if datay_rank.size(0) > 0: - # datay_ids = datay_rank[rank_train_mask].long().view(-1, 1) - # tstart_loss_calc = start_time(row_groups[0], rank) - - datay_ids = datay_rank[rank_train_mask].long() - - filtered_indices = torch.mul(datay_ids >= min_class, datay_ids < max_class).float() - indices = torch.nonzero(filtered_indices * torch.cuda.FloatTensor(datay_ids.size(), device=device).fill_(1)).squeeze() - - datay_ids = datay_rank[rank_train_mask].long().view(-1, 1) - datay_ids = datay_ids.index_select(0, indices) - datay_ids -= min_class - outputs_ids = outputs.index_select(0, indices) - - # classes = torch.gather(outputs[rank_train_mask], 1, datay_ids) - classes = torch.gather(outputs_ids, 1, datay_ids) - loss_calc = torch.sum(classes) - loss_calc_tens = torch.Tensor([loss_calc.item()]) - - rank_row_src = rank_row * proc_col - - # dist.reduce_multigpu([loss_calc], dst=rank_row_src, op=dist.reduce_op.SUM, group=row_groups[rank_row]) - # dist.broadcast_multigpu([loss_calc], src=rank_row_src, group=row_groups[rank_row]) - dist.reduce(loss_calc, dst=rank_row_src, op=dist.reduce_op.SUM, group=row_groups[rank_row]) - dist.broadcast(loss_calc, src=rank_row_src, group=row_groups[rank_row]) - - vertex_train_count = (data.train_mask.size(0) - (data.train_mask == 0).sum(dim=0)) - loss_calc = -loss_calc / vertex_train_count - - # loss_calc_time[run][rank] += stop_time(row_groups[0], rank, tstart_loss_calc) - - loss_calc.backward() - # print("loss_calc: " + str(loss_calc), flush=True) - # loss = F.nll_loss(outputs[rank_train_mask], datay_rank[rank_train_mask]) - # loss.backward() - # print("loss: " + str(loss), flush=True) - else: - fake_loss = (outputs * torch.cuda.FloatTensor(outputs.size(), device=device).fill_(0)).sum() - # fake_loss = (outputs * torch.zeros(outputs.size())).sum() - fake_loss.backward() - - optimizer.step() - - return outputs - -def test(outputs, data, vertex_count, rank): - logits, accs = outputs, [] - for _, mask in data('train_mask', 'val_mask', 'test_mask'): - pred = logits[mask].max(1)[1] - acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - accs.append(acc) - - if len(accs) != 3: - accs = accs + [0] * (3 - len(accs)) - - return accs - # logits, accs = outputs, [] - # datay_rank = torch.split(data.y, vertex_count)[rank] - # for _, mask in data('train_mask', 'val_mask', 'test_mask'): - # mask_rank = torch.split(mask, vertex_count)[rank] - # count = mask_rank.nonzero().size(0) - # if count > 0: - # pred = logits[mask_rank].max(1)[1] - # acc = pred.eq(datay_rank[mask_rank]).sum().item() / mask_rank.sum().item() - # # pred = logits[mask].max(1)[1] - # # acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - # else: - # acc = -1 - # accs.append(acc) - # return accs - - -# Split a COO into partitions of size n_per_proc -# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. -def split_coo(adj_matrix, node_count, n_per_proc, dim, size): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - vtx_indices = list(range(0, node_count, n_per_proc)) - vtx_indices = vtx_indices[:proc_row] - vtx_indices.append(node_count) - - am_partitions = [] - for i in range(len(vtx_indices) - 1): - am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] - am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] - am_part[dim] -= vtx_indices[i] - am_partitions.append(am_part) - - return am_partitions, vtx_indices - -# Normalize all elements according to KW's normalization rule -def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx): - if not normalization: - return adj_part - - adj_part = adj_part.coalesce() - deg = torch.histc(adj_matrix[0].double(), bins=node_count) - deg = deg.pow(-0.5) - - row_len = adj_part.size(0) - col_len = adj_part.size(1) - - dleft = torch.sparse_coo_tensor([np.arange(0, row_len).tolist(), - np.arange(0, row_len).tolist()], - deg[row_vtx:(row_vtx + row_len)].float(), - size=(row_len, row_len), - requires_grad=False, device=torch.device("cpu")) - - dright = torch.sparse_coo_tensor([np.arange(0, col_len).tolist(), - np.arange(0, col_len).tolist()], - deg[col_vtx:(col_vtx + col_len)].float(), - size=(col_len, col_len), - requires_grad=False, device=torch.device("cpu")) - # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) - ad_ind, ad_val = torch_sparse.spspmm(adj_part._indices(), adj_part._values(), - dright._indices(), dright._values(), - adj_part.size(0), adj_part.size(1), dright.size(1)) - - adj_part_ind, adj_part_val = torch_sparse.spspmm(dleft._indices(), dleft._values(), - ad_ind, ad_val, - dleft.size(0), dleft.size(1), adj_part.size(1)) - - adj_part = torch.sparse_coo_tensor(adj_part_ind, adj_part_val, - size=(adj_part.size(0), adj_part.size(1)), - requires_grad=False, device=torch.device("cpu")) - - return adj_part - -def proc_row_size(size): - return math.floor(math.sqrt(size)) - -def proc_col_size(size): - return math.floor(math.sqrt(size)) - -def twod_partition(rank, size, inputs, adj_matrix, data, features, classes, device): - node_count = inputs.size(0) - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - - inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) - - # n_per_proc = math.ceil(float(node_count) / proc_row) - n_per_proc = node_count // proc_row - - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - - am_partitions = None - am_pbyp = None - - # Compute the adj_matrix and inputs partitions for this process - # TODO: Maybe I do want grad here. Unsure. - with torch.no_grad(): - # Column partitions - am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1, size) - - proc_node_count = vtx_indices[rank_col + 1] - vtx_indices[rank_col] - am_pbyp, _ = split_coo(am_partitions[rank_col], node_count, n_per_proc, 0, size) - for i in range(len(am_pbyp)): - if i == proc_row - 1: - last_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(last_node_count, proc_node_count), - requires_grad=False) - - am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank_col]) - else: - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(n_per_proc, proc_node_count), - requires_grad=False) - - am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank_col]) - - # input_rowparts = torch.split(inputs, math.ceil(float(inputs.size(0)) / proc_row), dim=0) - inputs_per_row = inputs.size(0) // proc_row - inputs_per_col = inputs.size(1) // proc_col - chunks_per_row = [] - chunks_per_col = [] - for i in range(proc_row): - if i == proc_row - 1: - chunks_per_row.append(inputs.size(0) - inputs_per_row * (proc_row - 1)) - else: - chunks_per_row.append(inputs_per_row) - for i in range(proc_col): - if i == proc_col - 1: - chunks_per_col.append(inputs.size(1) - inputs_per_col * (proc_col - 1)) - else: - chunks_per_col.append(inputs_per_col) - - # input_rowparts = torch.split(inputs, math.ceil(float(inputs.size(0)) / proc_row), dim=0) - input_rowparts = torch.split(inputs, chunks_per_row, dim=0) - input_partitions = [] - for i in input_rowparts: - # input_partitions.append(torch.split(i, math.ceil(float(inputs.size(1)) / proc_col), - # dim=1)) - input_partitions.append(torch.split(i, chunks_per_col, dim=1)) - - adj_matrix_loc = am_pbyp[rank_row] - inputs_loc = input_partitions[rank_row][rank_col] - - print(adj_matrix_loc.size(), flush=True) - print(inputs_loc.size(), flush=True) - return inputs_loc, adj_matrix_loc, am_pbyp - -def rank_to_devid(rank, acc_per_rank): - return rank % acc_per_rank - -def run(rank, size, inputs, adj_matrix, data, features, mid_layer, classes, device, acc_per_rank): - global comm_time - global comp_time - global epochs - global timing - global run - - best_val_acc = test_acc = 0 - outputs = None - - group = dist.new_group(list(range(size))) - row_groups, col_groups = get_proc_groups(rank, size, group) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - rank_t = rank_col * proc_row + rank_row - - if rank_row >= proc_row or rank_col >= proc_col: - return - - transpose_groups = [] - - for i in range(proc_row): - transpose_groups_row = [] - for j in range(proc_col): - local_rank = i * proc_col + j - local_rank_t = j * proc_row + i - if local_rank < local_rank_t: - transpose_groups_row.append(dist.new_group([local_rank, local_rank_t])) - else: - transpose_groups_row.append(None) - transpose_groups.append(transpose_groups_row) - - if rank < rank_t: - transpose_group = transpose_groups[rank_row][rank_col] - else: - transpose_group = transpose_groups[rank_col][rank_row] - - # adj_matrix_loc = torch.rand(node_count, n_per_proc) - # inputs_loc = torch.rand(n_per_proc, inputs.size(1)) - - for i in range(run_count): - run = i - torch.manual_seed(0) - weight1_nonleaf = torch.rand(features, mid_layer, requires_grad=True) - weight1_nonleaf = weight1_nonleaf.to(device) - weight1_nonleaf.retain_grad() - - weight2_nonleaf = torch.rand(mid_layer, classes, requires_grad=True) - weight2_nonleaf = weight2_nonleaf.to(device) - weight2_nonleaf.retain_grad() - - weight1 = Parameter(weight1_nonleaf) - weight2 = Parameter(weight2_nonleaf) - - optimizer = torch.optim.Adam([weight1, weight2], lr=0.01) - - inputs_loc, adj_matrix_loc, _ = twod_partition(rank, size, inputs, adj_matrix, data, features, - classes, device) - - adj_matrix_loc = adj_matrix_loc.coalesce() - - inputs_loc = inputs_loc.to(device) - adj_matrix_loc = adj_matrix_loc.to(device) - - print(f"rank: {rank} adj_matrix_loc.nnz: {adj_matrix_loc._nnz()}") - - total_time[i] = dict() - comp_time[i] = dict() - comm_time[i] = dict() - summa_sparse_bcast1[i] = dict() - summa_sparse_bcast1_words[i] = dict() - summa_sparse_bcast2_words[i] = dict() - summa_sparse_bcast2[i] = dict() - summa_sparse_bcast2_fwd[i] = dict() - summa_sparse_bcast2_bwd[i] = dict() - summa_bcast1[i] = dict() - summa_bcast2[i] = dict() - summa_sparse_comp[i] = dict() - summa_comp[i] = dict() - summa_loc_bcast[i] = dict() - fwd_time[i] = dict() - bwd_time[i] = dict() - transpose_time[i] = dict() - grad_weight_time[i] = dict() - loss_calc_time[i] = dict() - summa_sparse_time[i] = dict() - summa_time[i] = dict() - summa_loc_time[i] = dict() - - total_time[i][rank] = 0.0 - comp_time[i][rank] = 0.0 - comm_time[i][rank] = 0.0 - summa_sparse_bcast1[i][rank] = 0.0 - summa_sparse_bcast1_words[i][rank] = 0.0 - summa_sparse_bcast2_words[i][rank] = 0.0 - summa_sparse_bcast2[i][rank] = 0.0 - summa_sparse_bcast2_fwd[i][rank] = 0.0 - summa_sparse_bcast2_bwd[i][rank] = 0.0 - summa_bcast1[i][rank] = 0.0 - summa_bcast2[i][rank] = 0.0 - summa_sparse_comp[i][rank] = 0.0 - summa_comp[i][rank] = 0.0 - summa_loc_bcast[i][rank] = 0.0 - fwd_time[i][rank] = 0.0 - bwd_time[i][rank] = 0.0 - transpose_time[i][rank] = 0.0 - grad_weight_time[i][rank] = 0.0 - loss_calc_time[i][rank] = 0.0 - summa_sparse_time[i][rank] = 0.0 - summa_time[i][rank] = 0.0 - summa_loc_time[i][rank] = 0.0 - - # Do not time first epoch - # timing_on = timing == True - # timing = False - # outputs = train(inputs_loc, weight1, weight2, inputs.size(0), adj_matrix_loc, None, - # optimizer, data, rank, size, acc_per_rank, group, row_groups, - # col_groups, transpose_group) - # if timing_on: - # timing = True - - # # tstart = start_time(group, rank) - dist.barrier(group) - tstart = time.time() - - print(f"Starting training... rank {rank} run {i}", flush=True) - for epoch in range(0, epochs): - outputs = train(inputs_loc, weight1, weight2, inputs.size(0), adj_matrix_loc, None, - optimizer, data, rank, size, acc_per_rank, group, row_groups, - col_groups, transpose_group) - print("Epoch: {:03d}".format(epoch), flush=True) - - # dur = stop_time(group, rank, tstart) - tstop = time.time() - total_time[i][rank] = tstop - tstart - - # Get median runtime according to rank0 and print that run's breakdown - dist.barrier(group) - if rank == 0: - total_times_r0 = [] - for i in range(run_count): - total_times_r0.append(total_time[i][0]) - - print(f"total_times_r0: {total_times_r0}") - median_run_time = statistics.median(total_times_r0) - median_idx = total_times_r0.index(median_run_time) - median_idx = torch.cuda.LongTensor([median_idx]) - else: - median_idx = torch.cuda.LongTensor([0]) - - # dist.broadcast(median_idx, src=0, group=group) - median_idx = median_idx.item() - print(f"rank: {rank} median_idx: {median_idx}") - print(f"rank: {rank} Time: {total_time[median_idx][rank]}") - print(f"rank: {rank} comm_time: {comm_time[median_idx][rank]}") - print(f"rank: {rank} comp_time: {comp_time[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_comp: {summa_sparse_comp[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_bcast1: {summa_sparse_bcast1[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_bcast1_words: {summa_sparse_bcast1_words[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_bcast2: {summa_sparse_bcast2[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_bcast2_fwd: {summa_sparse_bcast2_fwd[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_bcast2_bwd: {summa_sparse_bcast2_bwd[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_bcast2_words: {summa_sparse_bcast2_words[median_idx][rank]}") - print(f"rank: {rank} summa_comp: {summa_comp[median_idx][rank]}") - print(f"rank: {rank} summa_bcast1: {summa_bcast1[median_idx][rank]}") - print(f"rank: {rank} summa_bcast2: {summa_bcast2[median_idx][rank]}") - print(f"rank: {rank} summa_loc_bcast: {summa_loc_bcast[median_idx][rank]}") - print(f"rank: {rank} transpose_time: {transpose_time[median_idx][rank]}") - print(f"rank: {rank} grad_weight_time: {grad_weight_time[median_idx][rank]}") - print(f"rank: {rank} loss_calc_time: {loss_calc_time[median_idx][rank]}") - print(f"rank: {rank} summa_sparse_time: {summa_sparse_time[median_idx][rank]}") - print(f"rank: {rank} summa_time: {summa_time[median_idx][rank]}") - print(f"rank: {rank} summa_loc_time: {summa_loc_time[median_idx][rank]}") - print(f"rank: {rank} {outputs}") - - # All-gather outputs to test accuracy - if accuracy: - # All-gather across process row - output_parts_row = [] - width_per_proc = classes // proc_col - for i in range(proc_col): - output_parts_row.append(torch.cuda.FloatTensor(outputs.size(0), classes - width_per_proc * (proc_col - 1))) - - if outputs.size(1) != classes - width_per_proc * (proc_col - 1): - pad_col = (classes - width_per_proc * (proc_col - 1)) - outputs.size(1) - outputs = torch.cat((outputs, torch.cuda.FloatTensor(outputs.size(0), pad_col, device=device)), dim=1) - - dist.all_gather(output_parts_row, outputs, group=row_groups[rank_row]) - for i in range(proc_col - 1): - output_parts_row[i] = output_parts_row[i][:,:width_per_proc] - - outputs_row = torch.cat(output_parts_row, dim=1) - - # All-gather across process col - output_parts_col = [] - height_per_proc = inputs.size(0) // proc_row - for i in range(proc_row): - output_parts_col.append(torch.cuda.FloatTensor(inputs.size(0) - height_per_proc * (proc_row - 1), classes)) - - if outputs_row.size(0) != inputs.size(0) - height_per_proc * (proc_row - 1): - pad_row = (inputs.size(0) - height_per_proc * (proc_col - 1)) - outputs_row.size(0) - outputs_row = torch.cat((outputs_row, torch.cuda.FloatTensor(pad_row, classes, device=device)), dim=0) - - dist.all_gather(output_parts_col, outputs_row, group=col_groups[rank_col]) - for i in range(proc_row - 1): - output_parts_col[i] = output_parts_col[i][:height_per_proc,:] - - outputs = torch.cat(output_parts_col, dim=0) - - train_acc, val_acc, tmp_test_acc = test(outputs, data, inputs.size(0), rank) - if val_acc > best_val_acc: - best_val_acc = val_acc - test_acc = tmp_test_acc - log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}' - - print(log.format(200, train_acc, best_val_acc, test_acc)) - return outputs - -def init_process(rank, size, inputs, adj_matrix, data, features, mid_layer, classes, device, - outputs, acc_per_rank, fn): - - run_outputs = fn(rank, size, inputs, adj_matrix, data, features, mid_layer, classes, device, - acc_per_rank) - if outputs is not None: - outputs[rank] = run_outputs.detach() - -def main(): - # graphname = 'Reddit' - global graphname - global mid_layer - - path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - - # mid_layer = 16 - if graphname == 'Cora': - dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures()) - data = dataset[0] - num_features = dataset.num_features - num_classes = dataset.num_classes - elif graphname == 'Reddit': - dataset = Reddit(path, T.NormalizeFeatures()) - data = dataset[0] - num_features = dataset.num_features - # num_classes = dataset.num_classes + 9 - num_classes = dataset.num_classes - print(f"before edge_index: {data.edge_index.size()}") - elif graphname == 'Amazon': - # edge_index = torch.load(path + "/processed/amazon_graph.pt") - # edge_index = torch.load("/gpfs/alpine/bif115/scratch/alokt/Amazon/processed/amazon_graph_random.pt") - # edge_index = torch.load("/gpfs/alpine/bif115/scratch/alokt/Amazon/processed/amazon_large_randomized.pt") - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/Amazon/processed/data.pt") - print(f"Done loading coo", flush=True) - # edge_index = edge_index.t_() - # n = 9430086 - # n = 9430088 - n = 14249639 - num_features = 300 - num_classes = 24 - # mid_layer = 24 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1) - data.train_mask = torch.ones(n).long() - print(f"before edge_index: {edge_index.size()}") - elif graphname == 'subgraph5': - path = "/gpfs/alpine/bif115/scratch/alokt/HipMCL/" - edge_index = torch.load(path + "/processed/subgraph5_graph.pt") - n = 2186385 - num_features = 128 - # mid_layer = 64 - num_classes = 256 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1) - data.train_mask = torch.ones(n).long() - elif graphname == 'subgraph3': - # path = "/gpfs/alpine/bif115/scratch/alokt/HipMCL/" - # edge_index = torch.load(path + "/processed/subgraph3_graph.pt") - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/subgraph3/processed/data.pt") - print(f"Done loading coo", flush=True) - n = 8745542 - num_features = 128 - # mid_layer = 512 - # mid_layer = 64 - num_classes = 256 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1) - data.train_mask = torch.ones(n).long() - - if download: - exit() - - if "OMPI_COMM_WORLD_RANK" in os.environ.keys(): - os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] - - # Initialize distributed environment with SLURM - if "SLURM_PROCID" in os.environ.keys(): - os.environ["RANK"] = os.environ["SLURM_PROCID"] - - if "SLURM_NTASKS" in os.environ.keys(): - os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] - - if "MASTER_ADDR" not in os.environ.keys(): - os.environ["MASTER_ADDR"] = "127.0.0.1" - - os.environ["MASTER_PORT"] = "1234" - dist.init_process_group(backend='nccl') - # dist.init_process_group('gloo', init_method='env://') - rank = dist.get_rank() - size = dist.get_world_size() - - mp.set_start_method('spawn', force=True) - # device = torch.device('cpu') - devid = rank_to_devid(rank, acc_per_rank) - - device = torch.device('cuda:{}'.format(devid)) - print("device: " + str(device), flush=True) - torch.cuda.set_device(device) - curr_devid = torch.cuda.current_device() - devcount = torch.cuda.device_count() - print(f"devid: {curr_devid} {devcount}", flush=True) - - if graphname == "Amazon": - # edge_index = edge_index.to(device) - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - # inputs = inputs.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - elif graphname == "subgraph5": - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - elif graphname == "subgraph3": - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - else: - data = data.to(device) - data.x.requires_grad = True - inputs = data.x.to(device) - inputs.requires_grad = True - data.y = data.y.to(device) - - edge_index = data.edge_index - - if normalization: - adj_matrix, _ = add_remaining_self_loops(edge_index) - else: - adj_matrix = edge_index - - outputs = None - print("Processes: " + str(size), flush=True) - - # init_process(rank, size, inputs, adj_matrix, data, dataset.num_features, mid_layer, dataset.num_classes, - init_process(rank, size, inputs, adj_matrix, data, num_features, mid_layer, num_classes, - device, outputs, acc_per_rank, run) - - if outputs is not None: - return outputs[0] - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument("--local_rank", type=int) - parser.add_argument("--accperrank", type=int) - parser.add_argument("--epochs", type=int) - parser.add_argument("--graphname", type=str) - parser.add_argument("--timing", type=str) - parser.add_argument("--midlayer", type=int) - parser.add_argument("--runcount", type=int) - parser.add_argument("--normalization", type=str) - parser.add_argument("--activations", type=str) - parser.add_argument("--accuracy", type=str) - parser.add_argument("--download", type=bool) - args = parser.parse_args() - print(args) - - acc_per_rank = args.accperrank - if acc_per_rank is None: - acc_per_rank = 1 - - epochs = args.epochs - graphname = args.graphname - timing = args.timing == "True" - mid_layer = args.midlayer - run_count = args.runcount - normalization = args.normalization == "True" - activations = args.activations == "True" - accuracy = args.accuracy == "True" - download = args.download - - if not download: - if (epochs is None) or (graphname is None) or (timing is None) or (mid_layer is None) or (run_count is None): - print(f"Error: missing argument {epochs} {graphname} {timing} {mid_layer}") - exit() - - print(f"Arguments: epochs: {epochs} graph: {graphname} timing: {timing} mid: {mid_layer} norm: {normalization} act: {activations} acc: {accuracy}") - - print(main()) diff --git a/gcn_distr_3d.py b/gcn_distr_3d.py deleted file mode 100644 index 20bd16b..0000000 --- a/gcn_distr_3d.py +++ /dev/null @@ -1,1807 +0,0 @@ -import os -import os.path as osp -import argparse - -import math - -import torch -import torch_sparse -import torch.distributed as dist - -from torch_geometric.data import Data, Dataset -from torch_geometric.datasets import Planetoid, PPI -from reddit import Reddit -from torch_geometric.nn import GCNConv, ChebConv # noqa -from torch_geometric.utils import ( - add_remaining_self_loops, - to_dense_adj, - dense_to_sparse, - to_scipy_sparse_matrix -) -import torch_geometric.transforms as T - -import torch.multiprocessing as mp - -from torch.multiprocessing import Manager, Process - -from torch.nn import Parameter -import torch.nn.functional as F - -from torch_scatter import scatter_add - -import socket -import time -import numpy as np - -from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu - -comp_time = 0.0 -comm_time = 0.0 -summa_sparse_bcast1 = 0.0 -summa_sparse_bcast1_words = 0 -summa_sparse_bcast2_words = 0 -summa_sparse_bcast2 = 0.0 -summa_sparse_bcast2_fwd = 0.0 -summa_sparse_bcast2_bwd = 0.0 -summa_sparse_reduce = 0.0 -summa_bcast1 = 0.0 -summa_bcast2 = 0.0 -summa_reduce = 0.0 -summa_sparse_comp = 0.0 -summa_comp = 0.0 -summa_loc_bcast = 0.0 -fwd_time = 0.0 -bwd_time = 0.0 -transpose_time = 0.0 -grad_weight_time = 0.0 -loss_calc_time = 0.0 -summa_sparse_time = 0.0 -summa_time = 0.0 -summa_loc_time = 0.0 - -epochs = 0 -graphname = "" -mid_layer = 0 -timing = False -normalization = False -no_occur_val = 42.1234 - -def sync_and_sleep(rank, device): - torch.cuda.synchronize(device=device) - print(f"rank: {rank} memory_allocated: {torch.cuda.memory_allocated(device)}", flush=True) - # print(f"Sleeping rank {rank}", flush=True) - # time.sleep(20) - # print(f"Done sleeping rank {rank}", flush=True) - -def normalize(adj_matrix): - adj_matrix = adj_matrix + torch.eye(adj_matrix.size(0)) - d = torch.sum(adj_matrix, dim=1) - d = torch.rsqrt(d) - d = torch.diag(d) - return torch.mm(d, torch.mm(adj_matrix, d)) - -def start_time(group, rank): - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - if not timing: - return 0.0 - # dist.barrier(group) - torch.cuda.synchronize(device=device) - tstart = 0.0 - if rank == 0: - tstart = time.time() - return tstart - -def stop_time(group, rank, tstart): - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - if not timing: - return 0.0 - # dist.barrier(group) - torch.cuda.synchronize(device=device) - tstop = 0.0 - if rank == 0: - tstop = time.time() - return tstop - tstart - -def transpose(mat, rank, height, width, height_c, width_c, size, acc_per_rank, c_groups, transpose_group): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - # device = torch.device('cpu') - - rank_t = rank_col * proc_col * proc_c + rank_row * proc_c + rank_c - - width_per_proc_c = width_c // proc_c - mat_c_recv_width = width_c - width_per_proc_c * (proc_c - 1) - - if mat.size(1) != mat_c_recv_width: - pad_col = mat_c_recv_width - mat.size(1) - # mat = torch.cat((mat, torch.FloatTensor(mat.size(0), pad_col, device=device).fill_(no_occur_val)), - mat = torch.cat((mat, torch.cuda.FloatTensor(mat.size(0), pad_col, device=device).fill_(no_occur_val)), - dim=1) - - mat_c_recv = [] - for i in range(proc_c): - # mat_c_recv.append(torch.FloatTensor(mat.size(), device=device)) - mat_c_recv.append(torch.cuda.FloatTensor(mat.size(), device=device)) - - dist.all_gather(mat_c_recv, mat, group=c_groups[int(rank // proc_c)]) - for i in range(proc_c): - mat_c_recv[i] = mat_c_recv[i].t_() - mat_c_recv[i] = mat_c_recv[i][(mat_c_recv[i][:, 0] != no_occur_val).nonzero().squeeze(1)] - mat_c_recv[i] = mat_c_recv[i].t_() - - mat = torch.cat(mat_c_recv, dim=1) - mat.t_() - - if rank_col == proc_row - 1: - width_recv = height - (height // proc_row) * (proc_row - 1) - else: - width_recv = height // proc_row - - if rank_row == proc_col - 1: - height_recv = width - (width // proc_col) * (proc_col - 1) - else: - height_recv = width // proc_col - - if rank_row == rank_col: - mat_recv = mat - else: - mat_recv = torch.cuda.FloatTensor(height_recv, width_recv, device=device) - mat_recvs = [mat.contiguous(), mat_recv] - - if rank < rank_t: - dist.broadcast(mat_recvs[0], src=rank, group=transpose_group) - dist.broadcast(mat_recvs[1], src=rank_t, group=transpose_group) - else: - dist.broadcast(mat_recvs[1], src=rank_t, group=transpose_group) - dist.broadcast(mat_recvs[0], src=rank, group=transpose_group) - - chunk_sizes_col = [] - chunk_len = width_recv // proc_c - for i in range(proc_c): - if i == proc_c - 1: - chunk_sizes_col.append(width_recv - chunk_len * (proc_c - 1)) - else: - chunk_sizes_col.append(chunk_len) - mat_split = torch.split(mat_recv, chunk_sizes_col, dim=1) - - del mat_recv - - return mat_split[rank_c] - -def split3dspmm_dense(adj_matrix, inputs, rank, row, col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, - height, middim, width): - - global comm_time - global comp_time - - global summa_bcast1 - global summa_bcast2 - global summa_reduce - - global summa_comp - global summa_sparse_time - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - # height_per_proc = math.ceil(float(height) / proc_row) - # width_per_proc = math.ceil(float(width) / proc_col) - - # # TODO: Not sure how to handle this w/o square grid - # middim_per_proc = math.ceil(float(middim) / proc_col) - height_per_proc = height // proc_row - width_per_proc = width // proc_col - - # TODO: Not sure how to handle this w/o square grid - middim_per_proc = middim // (proc_col * proc_c) - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - # device = torch.device('cpu') - - if row == proc_row - 1: - # height_per_proc -= proc_row * height_per_proc - height - height_per_proc = height - height_per_proc * (proc_row - 1) - - if col == proc_col - 1: - # width_per_proc -= proc_col * width_per_proc - width - width_per_proc = width - width_per_proc * (proc_col - 1) - - height_per_proc_c = adj_matrix.size(0) // proc_c - if rank_c == proc_c - 1: - height_per_proc_c = adj_matrix.size(0) - height_per_proc_c * (proc_c - 1) - - z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - # z_loc = torch.cuda.FloatTensor(height_per_proc_c, width_per_proc, device=device).fill_(0) - # width_per_proc_c = inputs.size(1) // proc_c - # if rank_c == proc_c - 1: - # width_per_proc_c = inputs.size(1) - width_per_proc_c * (proc_c - 1) - - # z_loc = torch.zeros(height_per_proc, width_per_proc_c, device=device) - - # chunk_sizes_col = [] - # chunk_len = inputs.size(1) // proc_c - # for i in range(proc_c): - # if i == proc_c - 1: - # chunk_sizes_col.append(inputs.size(1) - chunk_len * (proc_c - 1)) - # else: - # chunk_sizes_col.append(chunk_len) - - chunk_sizes_row = [] - chunk_len = adj_matrix.size(0) // proc_c - for i in range(proc_c): - if i == proc_c - 1: - chunk_sizes_row.append(adj_matrix.size(0) - chunk_len * (proc_c - 1)) - else: - chunk_sizes_row.append(chunk_len) - - for k in range(proc_col): - - row_src_rank = row * (proc_col * proc_c) + rank_c + k * proc_c - col_src_rank = col * proc_row + rank_c + k * proc_c * proc_row - - middim_per_col = middim // proc_col - if k == proc_col - 1: - middim_per_col = middim - middim_per_col * (proc_col - 1) - - middim_per_proc = middim_per_col // proc_c - if rank_c == proc_c - 1: - middim_per_proc = middim_per_col - middim_per_proc * (proc_c - 1) - - - if row_src_rank == rank: - acol = adj_matrix - else: - acol = torch.cuda.FloatTensor(height_per_proc, middim_per_proc, device=device).fill_(0) - - tstart = start_time(row_groups[row][rank_c], rank) - - acol = acol.contiguous() - dist.broadcast(acol, row_src_rank, row_groups[row][rank_c]) - - dur = stop_time(row_groups[row][rank_c], rank, tstart) - comm_time += dur - summa_bcast1 += dur - - if col_src_rank == rank: - brow = inputs - else: - brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=device).fill_(0) - # brow = torch.FloatTensor(middim_per_proc, width_per_proc, device=device).fill_(0) - - tstart = start_time(row_groups[0][0], rank) - - brow = brow.contiguous() - dist.broadcast(brow, col_src_rank, col_groups[col][rank_c]) - - dur = stop_time(row_groups[0][0], rank, tstart) - comm_time += dur - summa_bcast2 += dur - - tstart = start_time(row_groups[0][0], rank) - - z_loc += torch.mm(acol, brow) - - dur = stop_time(row_groups[0][0], rank, tstart) - comp_time += dur - summa_comp += dur - - del acol - del brow - - # tstart = start_time(row_groups[0][0], rank) - tstart = start_time(c_groups[0], rank) - - dist.all_reduce(z_loc, group=c_groups[int(rank // proc_c)]) - z_loc = torch.split(z_loc, chunk_sizes_row, dim=0) - z_loc = z_loc[rank_c].contiguous() - - - # dur = stop_time(row_groups[0][0], rank, tstart) - dur = stop_time(c_groups[0], rank, tstart) - comm_time += dur - summa_reduce += dur - # dist.all_gather(z_tmp_recv, z_tmp, group=c_groups[int(rank // proc_c)]) - # z_tmp_recv[rank_c] = z_tmp - - # z_tmp = z_tmp_recv[0] - # for i in range(1, proc_c): - # z_tmp += z_tmp_recv[i] - - - # z_tmp_recv = [] - # for i in range(proc_c): - # z_tmp_recv.append(torch.FloatTensor(z_tmp.size(), device=device)) - - return z_loc - -def split3dspmm_sparse(adj_matrix, inputs, rank, row, col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, - height, middim, width): - - global comm_time - global comp_time - - global summa_sparse_bcast1 - global summa_sparse_bcast2 - global summa_sparse_reduce - - global summa_sparse_bcast1_words - global summa_sparse_bcast2_words - - global summa_sparse_comp - global summa_sparse_time - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - # height_per_proc = math.ceil(float(height) / proc_row) - # width_per_proc = math.ceil(float(width) / proc_col) - - # # TODO: Not sure how to handle this w/o square grid - # middim_per_proc = math.ceil(float(middim) / proc_col) - height_per_proc = height // proc_row - width_per_proc = width // proc_col - - # TODO: Not sure how to handle this w/o square grid - middim_per_proc = middim // (proc_col * proc_c) - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - # device = torch.device('cpu') - - if row == proc_row - 1: - # height_per_proc -= proc_row * height_per_proc - height - height_per_proc = height - height_per_proc * (proc_row - 1) - - if col == proc_col - 1: - # width_per_proc -= proc_col * width_per_proc - width - width_per_proc = width - width_per_proc * (proc_col - 1) - - width_per_proc_c = inputs.size(1) // proc_c - if rank_c == proc_c - 1: - width_per_proc_c = inputs.size(1) - width_per_proc_c * (proc_c - 1) - - # z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc_c, device=device).fill_(0) - z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - - chunk_sizes_col = [] - chunk_len = inputs.size(1) // proc_c - for i in range(proc_c): - if i == proc_c - 1: - chunk_sizes_col.append(inputs.size(1) - chunk_len * (proc_c - 1)) - else: - chunk_sizes_col.append(chunk_len) - - for k in range(proc_col): - - row_src_rank = row * (proc_col * proc_c) + rank_c + k * proc_c - col_src_rank = col * proc_row + rank_c + k * proc_c * proc_row - - # if k == proc_col * proc_c - 1: - # # middim_per_proc -= proc_col * middim_per_proc - middim - # middim_per_proc = middim - middim_per_proc * (proc_col * proc_c - 1) - middim_per_col = middim // proc_col - if k == proc_col - 1: - middim_per_col = middim - middim_per_col * (proc_col - 1) - - middim_per_proc = middim_per_col // proc_c - if rank_c == proc_c - 1: - middim_per_proc = middim_per_col - middim_per_proc * (proc_c - 1) - - if row_src_rank == rank: - # acol = adj_matrix.clone() - acol_indices_len = torch.cuda.LongTensor( - # acol_indices_len = torch.LongTensor( - [adj_matrix.indices().contiguous()[0].size(0)], - device=device) - acol_values_len = torch.cuda.LongTensor([adj_matrix.values().contiguous().size(0)], - # acol_values_len = torch.LongTensor([adj_matrix.values().contiguous().size(0)], - device=device) - else: - # acol = torch.sparse.FloatTensor(height_per_proc, middim_per_proc) - acol_indices_len = torch.cuda.LongTensor([0], device=device) - acol_values_len = torch.cuda.LongTensor([0], device=device) - - # acol_indices_len = torch.LongTensor([0], device=device) - # acol_values_len = torch.LongTensor([0], device=device) - - dist.broadcast(acol_indices_len, row_src_rank, row_groups[row][rank_c]) - - acol_indices_len = acol_indices_len.item() # nnz - acol_values_len = acol_indices_len - - if row_src_rank == rank: - acol_indices = adj_matrix.indices().contiguous().long() - acol_values = adj_matrix.values().contiguous().float() - else: - acol_indices = torch.cuda.LongTensor(2, acol_indices_len, device=device).fill_(0) - acol_values = torch.cuda.FloatTensor(acol_values_len, device=device).fill_(0) - # acol_indices = torch.LongTensor(2, acol_indices_len, device=device).fill_(0) - # acol_values = torch.FloatTensor(acol_values_len, device=device).fill_(0) - - - acol = torch.cat((acol_indices.float(), acol_values.unsqueeze(0)), dim=0) - - tstart = start_time(row_groups[row][rank_c], rank) - - dist.broadcast(acol.contiguous(), row_src_rank, row_groups[row][rank_c]) - - dur = stop_time(row_groups[row][rank_c], rank, tstart) - comm_time += dur - summa_sparse_bcast1 += dur - if rank == 0: - summa_sparse_bcast1_words += 3 * acol_values_len - - acol_indices = acol[:2].long() - acol_values = acol[2].squeeze(0) - - if row_src_rank == rank: - acol = adj_matrix - else: - acol = sparse_coo_tensor_gpu(acol_indices, acol_values, - torch.Size([height_per_proc, middim_per_proc])) - # acol = torch.sparse_coo_tensor(acol_indices, acol_values, - # torch.Size([height_per_proc, middim_per_proc]), - # requires_grad=False) - - - if col_src_rank == rank: - brow = inputs - else: - brow = torch.cuda.FloatTensor(middim_per_proc, width_per_proc, device=device) - # brow = torch.FloatTensor(middim_per_proc, width_per_proc, device=device) - - - tstart = start_time(row_groups[0][0], rank) - - brow = brow.contiguous() - dist.broadcast(brow, col_src_rank, col_groups[col][rank_c]) - - dur = stop_time(row_groups[0][0], rank, tstart) - comm_time += dur - summa_sparse_bcast2 += dur - if rank == 0: - summa_sparse_bcast2_words += brow.size(0) * brow.size(1) - - tstart = start_time(row_groups[0][0], rank) - - # z_tmp = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - spmm_gpu(acol_indices[0].int(), acol_indices[1].int(), acol_values, - height_per_proc, middim_per_proc, brow, z_loc) - # z_loc += torch.sparse.mm(acol, brow) - - dur = stop_time(row_groups[0][0], rank, tstart) - comp_time += dur - summa_sparse_comp += dur - - # del acol - # del brow - - # tstart = start_time(row_groups[0][0], rank) - z_loc = z_loc.contiguous() - tstart = start_time(c_groups[0], rank) - - dist.all_reduce(z_loc, group=c_groups[int(rank // proc_c)]) - z_loc = torch.split(z_loc, chunk_sizes_col, dim=1) - z_loc = z_loc[rank_c].contiguous() - - # dur = stop_time(row_groups[0][0], rank, tstart) - dur = stop_time(c_groups[0], rank, tstart) - comm_time += dur - summa_sparse_reduce += dur - - return z_loc, chunk_sizes_col - - -def split3dspmm_loc(mata, matb, rank, row, col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, - height, middim, width): - - - global comm_time - global comp_time - - global summa_bcast1 - global summa_bcast2 - global summa_reduce - - global summa_comp - global summa_sparse_time - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - # # TODO: Not sure how to handle this w/o square grid - height_per_proc = height // proc_row - width_per_proc = width // proc_col - # TODO: Not sure how to handle this w/o square grid - middim_per_proc = middim // (proc_col * proc_c) - - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - # device = torch.device('cpu') - - if row == proc_row - 1: - # height_per_proc -= proc_row * height_per_proc - height - height_per_proc = height - height_per_proc * (proc_row - 1) - - # if col == proc_col - 1: - # width_per_proc -= proc_col * width_per_proc - width - - width_per_proc = matb[rank].size(1) - - height_per_proc_c = mata.size(0) // proc_c - if rank_c == proc_c - 1: - height_per_proc_c = mata.size(0) - height_per_proc_c * (proc_c - 1) - - z_loc = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - # z_tmp = torch.cuda.FloatTensor(height_per_proc, width_per_proc, device=device).fill_(0) - # z_loc = torch.cuda.FloatTensor(height_per_proc_c, width_per_proc, device=device).fill_(0) - - chunk_sizes_row = [] - chunk_len = mata.size(0) // proc_c - for i in range(proc_c): - if i == proc_c - 1: - chunk_sizes_row.append(mata.size(0) - chunk_len * (proc_c - 1)) - else: - chunk_sizes_row.append(chunk_len) - - for k in range(proc_col): - - row_src_rank = row * (proc_col * proc_c) + rank_c + k * proc_c - col_src_rank = col * proc_row + rank_c + k * proc_c * proc_row - - middim_per_col = middim // proc_col - if k == proc_col - 1: - middim_per_col = middim - middim_per_col * (proc_col - 1) - - middim_per_proc = middim_per_col // proc_c - if rank_c == proc_c - 1: - middim_per_proc = middim_per_col - middim_per_proc * (proc_c - 1) - - if row_src_rank == rank: - acol = mata - else: - acol = torch.cuda.FloatTensor(height_per_proc, matb[col_src_rank].size(0), - device=device) - # acol = torch.FloatTensor(height_per_proc, matb[col_src_rank].size(0), - # device=device) - - tstart = start_time(row_groups[row][rank_c], rank) - - dist.broadcast(acol.contiguous(), row_src_rank, row_groups[row][rank_c]) - - dur = stop_time(row_groups[row][rank_c], rank, tstart) - comm_time += dur - summa_bcast1 += dur - - # if col_src_rank == rank: - # brow = matb.clone() - # else: - # brow = torch.FloatTensor(middim_per_proc, width_per_proc) - - # dist.broadcast(brow, col_src_rank, col_groups[col]) - - brow = matb[col_src_rank] - - tstart = start_time(row_groups[0][0], rank) - - z_tmp = torch.mm(acol, brow) - - dur = stop_time(row_groups[0][0], rank, tstart) - summa_comp += dur - comp_time += dur - - # tstart = start_time(row_groups[0][0], rank) - tstart = start_time(c_groups[0], rank) - - dist.all_reduce(z_tmp, group=c_groups[int(rank // proc_c)]) - - z_loc += z_tmp - - # dur = stop_time(row_groups[0][0], rank, tstart) - dur = stop_time(c_groups[0], rank, tstart) - summa_reduce += dur - comm_time += dur - - # del acol - # del z_tmp - - # z_tmp_recv = [] - # for i in range(proc_c): - # z_tmp_recv.append(torch.FloatTensor(z_tmp.size(), device=device)) - - # del mata - # del z_loc - # return z_loc_ret - - z_tmp = torch.split(z_loc, chunk_sizes_row, dim=0) - z_loc = z_tmp[rank_c].clone() - return z_loc - -def get_proc_groups(rank, size, group): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - row_groups = [] - col_groups = [] - c_groups = [] - - row_procs = [] - col_procs = [] - c_procs = [] - - for i in range(proc_row): - row_groups_c = [] - row_procs_c = [] - for j in range(proc_c): - proc_start = i * proc_col * proc_c + j - proc_end = (i + 1) * proc_col * proc_c + j - row_groups_c.append(dist.new_group(list(range(proc_start, proc_end, proc_c)))) - row_procs_c.append(list(range(proc_start, proc_end, proc_c))) - row_groups.append(row_groups_c) - row_procs.append(row_procs_c) - - for i in range(proc_col): - col_groups_c =[] - col_procs_c =[] - for j in range(proc_c): - proc_start = i * proc_c + j - proc_end = proc_row * proc_col * proc_c + i * proc_c + j - col_groups_c.append(dist.new_group(list(range(proc_start, proc_end, proc_c * proc_col)))) - col_procs_c.append(list(range(proc_start, proc_end, proc_c * proc_col))) - col_groups.append(col_groups_c) - col_procs.append(col_procs_c) - - for i in range(0, size, proc_c): - c_groups.append(dist.new_group(list(range(i, i + proc_c)))) - c_procs.append(list(range(i, i + proc_c))) - - return row_groups, col_groups, c_groups - -def dist_log_softmax(z, rank, size, acc_per_rank, group): - torch.set_printoptions(edgeitems=4) - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - maxes = torch.max(z, dim=1, keepdim=True)[0] - maxes_recv = [] - for i in range(proc_col): - maxes_recv.append(torch.cuda.FloatTensor(maxes.size(), device=device)) - - # dist.all_reduce(maxes, op=dist.reduce_op.MAX, group=group) - dist.all_gather(maxes_recv, maxes, group=group) - maxes_recv[rank_col] = maxes - maxes = torch.max(torch.cat(maxes_recv, dim=1), dim=1, keepdim=True)[0] - - h = torch.exp(z - maxes) - sm_sum = torch.sum(h, dim=1, keepdim=True) - - sm_sum_recv = [] - for i in range(proc_col): - sm_sum_recv.append(torch.cuda.FloatTensor(sm_sum.size(), device=device)) - - # dist.all_reduce(sm_sum, op=dist.reduce_op.SUM, group=group) - dist.all_gather(sm_sum_recv, sm_sum, group=group) - sm_sum_recv[rank_col] = sm_sum - sm_sum = torch.sum(torch.cat(sm_sum_recv, dim=1), dim=1, keepdim=True) - sm_sum = torch.log(sm_sum) - h = z - maxes - sm_sum - return h - -def dist_log_softmax2(z, rank, size, acc_per_rank, group): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - rank_row = int(rank / proc_col) - rank_col = rank % proc_col - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - maxes = torch.max(z, dim=1, keepdim=True)[0] - maxes_recv = [] - for i in range(proc_col): - maxes_recv.append(torch.cuda.FloatTensor(maxes.size(), device=device)) - - # dist.all_reduce(maxes, op=dist.reduce_op.MAX, group=group) - dist.all_gather(maxes_recv, maxes, group=group) - for i in range(proc_col): - maxes_recv[rank_col].requires_grad = True - maxes_recv[rank_col] = maxes - maxes = torch.max(torch.cat(maxes_recv, dim=1), dim=1, keepdim=True)[0] - - h = torch.exp(z - maxes) - # sm_sum = torch.sum(h, dim=1, keepdim=True) - - # sm_sum_recv = [] - # for i in range(proc_col): - # sm_sum_recv.append(torch.cuda.FloatTensor(sm_sum.size(), device=device)) - - # # dist.all_reduce(sm_sum, op=dist.reduce_op.SUM, group=group) - # dist.all_gather(sm_sum_recv, sm_sum, group=group) - # sm_sum_recv[rank_col] = sm_sum - # sm_sum = torch.sum(torch.cat(sm_sum_recv, dim=1), dim=1, keepdim=True) - # sm_sum = torch.log(sm_sum) - # h = z - maxes - sm_sum - return h - -class GCNFunc(torch.autograd.Function): - @staticmethod - def forward(ctx, inputs, weight, node_count, adj_matrix, am_partitions, rank, size, - acc_per_rank, group, row_groups, col_groups, transpose_group, c_groups, func): - # inputs: H - # adj_matrix: A - # weight: W - # func: sigma - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - # device = torch.device('cpu') - - ctx.save_for_backward(inputs, weight, adj_matrix) - ctx.node_count = node_count - ctx.rank = rank - ctx.size = size - ctx.acc_per_rank = acc_per_rank - ctx.group = group - ctx.row_groups = row_groups - ctx.col_groups = col_groups - ctx.c_groups = c_groups - ctx.transpose_group = transpose_group - - ctx.func = func - - adj_matrix_t = adj_matrix # Only true for undirected graphs - - # TODO: will need to change height argument when n % sqrt(P) != 0 and non-square grid - # z = summa_sparse(adj_matrix_t, inputs, rank, rank_row, rank_col, size, acc_per_rank, - z, chunk_sizes_loc = split3dspmm_sparse(adj_matrix_t, inputs, - # z, chunk_sizes_loc = split3dspmm_dense(adj_matrix_t, inputs, - rank, rank_row, rank_col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, - node_count, node_count, weight.size(0)) - - chunk_sizes_loc_tens = torch.cuda.LongTensor(chunk_sizes_loc) - chunk_sizes = [] - for i in range(proc_col): - chunk_sizes.append(torch.cuda.LongTensor(chunk_sizes_loc_tens.size())) - dist.all_gather(chunk_sizes, chunk_sizes_loc_tens, group=row_groups[rank_row][rank_c]) - chunk_sizes = torch.cat(chunk_sizes).tolist() - - chunk_sizes_row = [] - chunk_sizes_col = [] - weight_per_row = weight.size(0) // (proc_row * proc_c) - weight_per_col = weight.size(1) // proc_col - # for i in range(proc_row * proc_c): - # if i == proc_row * proc_c - 1: - # chunk_sizes_row.append(weight.size(0) - weight_per_row * (proc_row * proc_c - 1)) - # else: - # chunk_sizes_row.append(weight_per_row) - chunk_sizes_row = chunk_sizes - - for i in range(proc_col): - if i == proc_col - 1: - chunk_sizes_col.append(weight.size(1) - weight_per_col * (proc_col - 1)) - else: - chunk_sizes_col.append(weight_per_col) - - # weight_rows = torch.split(weight, math.ceil(float(weight.size(0)) / proc_row), dim=0) - weight_rows = torch.split(weight, chunk_sizes_row, dim=0) - weight_parts_tmp = [] - for i in weight_rows: - # weight_cols = torch.split(i, math.ceil(float(weight.size(1)) / proc_col), dim=1) - weight_cols = torch.split(i, chunk_sizes_col, dim=1) - weight_parts_tmp.extend(weight_cols) - - weight_parts = [None] * size - for i in range(proc_row * proc_c): - for j in range(proc_col): - rank_old = i * proc_col + j - - rank_new_row = i // proc_c - rank_new_c = i % proc_c - rank_new = rank_new_row * proc_col * proc_c + j * proc_c + rank_new_c - - if j == rank_col: - weight_parts[rank_new] = weight_parts_tmp[rank_old] - else: - weight_parts[rank_new] = None - - # z = torch.mm(z, weight) - z = split3dspmm_loc(z, weight_parts, rank, rank_row, rank_col, rank_c, size, acc_per_rank, row_groups, - col_groups, c_groups, node_count, weight.size(0), weight.size(1)) - - # z.requires_grad = True - ctx.z = z - - # Worry about activation later - # if func is F.log_softmax: - # h = dist_log_softmax(z, rank, size, acc_per_rank, row_groups[rank_row]) - # elif func is F.relu: - # h = func(z) - # else: - # h = z - - # fwd_time += dur - - return z - - @staticmethod - def backward(ctx, grad_output): - global summa_sparse_bcast2 - global summa_sparse_bcast2_bwd - global bwd_time - global transpose_time - global grad_weight_time - - inputs, weight, adj_matrix = ctx.saved_tensors - rank = ctx.rank - size = ctx.size - acc_per_rank = ctx.acc_per_rank - group = ctx.group - row_groups = ctx.row_groups - col_groups = ctx.col_groups - c_groups = ctx.c_groups - transpose_group = ctx.transpose_group - node_count = ctx.node_count - - func = ctx.func - z = ctx.z - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - # device = torch.device('cpu') - - # tstart = start_time(row_groups[0], rank) - - # Worry about activation later - # with torch.set_grad_enabled(True): - # if func is F.log_softmax: - # func_eval = dist_log_softmax2(z, rank, size, acc_per_rank, row_groups[rank_row]) - # elif func is F.relu: - # func_eval = func(z) - # else: - # func_eval = z - - # # sigmap = torch.autograd.grad(outputs=func_eval, inputs=z,grad_outputs=grad_output)[0] - # sigmap = torch.autograd.grad(outputs=func_eval, inputs=z,grad_outputs=grad_output)[0] - # print(f"rank: {rank} sigmap: {sigmap}", flush=True) - # grad_output = sigmap - - tmp_summa_sparse_bcast2 = summa_sparse_bcast2 - - # First backprop equation - # TODO: will need to change height argument when n % sqrt(P) != 0 and non-square grid - ag, chunk_sizes_loc = split3dspmm_sparse(adj_matrix, grad_output, - rank, rank_row, rank_col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, node_count, node_count, weight.t().size(0)) - - chunk_sizes_loc_tens = torch.cuda.LongTensor(chunk_sizes_loc) - chunk_sizes = [] - for i in range(proc_col): - chunk_sizes.append(torch.cuda.LongTensor(chunk_sizes_loc_tens.size())) - dist.all_gather(chunk_sizes, chunk_sizes_loc_tens, group=row_groups[rank_row][rank_c]) - chunk_sizes = torch.cat(chunk_sizes).tolist() - - # tstart_grad_weight = start_time(row_groups[0], rank) - chunk_sizes_row = [] - chunk_sizes_col = [] - weight_per_row = weight.t().size(0) // (proc_row * proc_c) - weight_per_col = weight.t().size(1) // proc_col - # for i in range(proc_row * proc_c): - # if i == proc_row * proc_c - 1: - # chunk_sizes_row.append(weight.t().size(0) - weight_per_row * (proc_row * proc_c - 1)) - # else: - # chunk_sizes_row.append(weight_per_row) - chunk_sizes_row = chunk_sizes - - for i in range(proc_col): - if i == proc_col - 1: - chunk_sizes_col.append(weight.t().size(1) - weight_per_col * (proc_col - 1)) - else: - chunk_sizes_col.append(weight_per_col) - # weight_rows = torch.split(weight.t(), math.ceil(float(weight.t().size(0)) / proc_row), - weight_rows = torch.split(weight.t(), chunk_sizes_row, dim=0) - - weight_parts_tmp = [] - for i in weight_rows: - # weight_cols = torch.split(i, math.ceil(float(weight.t().size(1)) / proc_col), dim=1) - weight_cols = torch.split(i, chunk_sizes_col, dim=1) - weight_parts_tmp.extend(weight_cols) - - weight_parts = [None] * size - for i in range(proc_row * proc_c): - for j in range(proc_col): - rank_old = i * proc_col + j - - rank_new_row = i // proc_c - rank_new_c = i % proc_c - rank_new = rank_new_row * proc_col * proc_c + j * proc_c + rank_new_c - - if rank_col == j: - weight_parts[rank_new] = weight_parts_tmp[rank_old] - else: - weight_parts[rank_new] = None - - # grad_input = torch.mm(ag, weight.t()) - grad_input = split3dspmm_loc(ag, weight_parts, rank, rank_row, rank_col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, - node_count, weight.t().size(0), weight.size(1)) - - # grad_weight_time += stop_time(row_groups[0], rank, tstart_grad_weight) - - # Second backprop equation (reuses the A * G^l computation) - # col_groups twice because of transpose - # TODO: will need to change height argument when n % sqrt(P) != 0 and non-square grid - - height_c = node_count // proc_row - if rank_row == proc_row - 1: - height_c = node_count - height_c * (proc_row - 1) - - width_c = weight.size(1) // proc_col - if rank_col == proc_col - 1: - width_c = weight.size(1) - width_c * (proc_col - 1) - - tstart_transpose = start_time(row_groups[0][0], rank) - - ag_t = transpose(ag, rank, node_count, weight.size(1), height_c, width_c, size, acc_per_rank, - c_groups, transpose_group) - - transpose_time += stop_time(row_groups[0][0], rank, tstart_transpose) - - # grad_weight = summa(inputs_t, ag, rank, rank_row, rank_col, size, acc_per_rank, row_groups, - # col_groups, weight.size(0), node_count, weight.size(1)) - grad_weight = split3dspmm_dense(ag_t, inputs, - rank, rank_row, rank_col, rank_c, size, acc_per_rank, - row_groups, col_groups, c_groups, weight.size(1), node_count, weight.size(0)) - - # tstart_grad_weight = start_time(row_groups[0], rank) - # Collect grad_weight's across processes - grad_weight_recv = [] - # max_row_chunk = max(chunk_sizes_col) #transpose - # max_col_chunk = max(chunk_sizes_row) - max_row_chunk = max(chunk_sizes_row) - max_col_chunk = max(chunk_sizes_col) - for i in range(size): - grad_weight_recv.append(torch.cuda.FloatTensor( - # grad_weight_recv.append(torch.FloatTensor( - max_row_chunk, - max_col_chunk, - device=device)) - - # pad_row = math.ceil(float(weight.size(0)) / proc_row) - grad_weight.size(0) - # pad_col = math.ceil(float(weight.size(1)) / proc_col) - grad_weight.size(1) - pad_row = max_row_chunk - grad_weight.size(0) - pad_col = max_col_chunk - grad_weight.size(1) - - # TODO: make this part less hacky - grad_weight = torch.cat((grad_weight, - torch.cuda.FloatTensor(pad_row, grad_weight.size(1), device=device).fill_(no_occur_val)), - # torch.FloatTensor(pad_row, grad_weight.size(1), device=device).fill_(no_occur_val)), - dim=0) - grad_weight = torch.cat((grad_weight, - torch.cuda.FloatTensor(grad_weight.size(0), pad_col, device=device).fill_(no_occur_val)), - # torch.FloatTensor(grad_weight.size(0), pad_col, device=device).fill_(no_occur_val)), - dim=1) - - dist.all_gather(grad_weight_recv, grad_weight) - - # for i in range(size): - # if rank == i: - # grad_weight_recv[i] = grad_weight - # dist.broadcast(grad_weight_recv[i], i, group) - # grad_weight_recv[0] = grad_weight - - for i in range(len(grad_weight_recv)): - grad_weight_recv[i] = grad_weight_recv[i][(grad_weight_recv[i][:, 0] != no_occur_val) - .nonzero().squeeze(1)] - - grad_weight_recv_t = grad_weight_recv[i].t() - grad_weight_recv_t = grad_weight_recv_t[(grad_weight_recv_t[:, 0] != no_occur_val) - .nonzero().squeeze(1)] - - grad_weight_recv[i] = grad_weight_recv_t.t() - - grad_weight_fin = torch.cuda.FloatTensor(device=device) - # grad_weight_fin = torch.FloatTensor(device=device) - for i in range(proc_row): - grad_weight_row = torch.cuda.FloatTensor(device=device) - # grad_weight_row = torch.FloatTensor(device=device) - for j in range(proc_col): - # grad_weight_col = torch.FloatTensor(device=device) - grad_weight_col = torch.cuda.FloatTensor(device=device) - for k in range(proc_c): - rank_wt = i * proc_row * proc_c + j * proc_c + k - grad_weight_col = torch.cat((grad_weight_col, grad_weight_recv[rank_wt]), dim=0) - grad_weight_row = torch.cat((grad_weight_row, grad_weight_col), dim=1) - grad_weight_fin = torch.cat((grad_weight_fin, grad_weight_row), dim=0) - - # dur = stop_time(row_groups[0], rank, tstart) - # bwd_time += dur - - # grad_weight_time += stop_time(row_groups[0], rank, tstart_grad_weight) - grad_weight_fin = grad_weight_fin.t() - - del ag - return grad_input, grad_weight_fin, None, None, None, None, None, None, None, None, None, None, None, None - -def train(inputs, weight1, weight2, node_count, adj_matrix, am_partitions, optimizer, data, rank, - size, acc_per_rank, group, row_groups, col_groups, transpose_group, c_groups): - - global loss_calc_time - - device = torch.device('cuda:{}'.format(rank_to_devid(rank, acc_per_rank))) - - outputs = GCNFunc.apply(inputs, weight1, node_count, adj_matrix, am_partitions, rank, size, - acc_per_rank, group, row_groups, col_groups, transpose_group, c_groups, - F.relu) - - outputs = GCNFunc.apply(outputs, weight2, node_count, adj_matrix, am_partitions, rank, size, - acc_per_rank, group, row_groups, col_groups, transpose_group, c_groups, - F.log_softmax) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - # device = torch.device('cpu') - - optimizer.zero_grad() - rank_train_mask = torch.split(data.train_mask, outputs.size(0), dim=0)[rank_row * proc_c + rank_c] - datay_rank = torch.split(data.y, outputs.size(0), dim=0)[rank_row * proc_c + rank_c] - - total_classes = weight2.size(1) - # class_per_rank = math.ceil(float(total_classes) / proc_col) - class_per_rank = total_classes // proc_col - - min_class = rank_col * class_per_rank - max_class = min((rank_col + 1) * class_per_rank, total_classes) - if rank_col == proc_col - 1: - max_class = total_classes - - # Note: bool type removes warnings, unsure of perf penalty - # loss = F.nll_loss(outputs[data.train_mask.bool()], data.y[data.train_mask.bool()]) - if list(datay_rank[rank_train_mask].size())[0] > 0: - # if datay_rank.size(0) > 0: - # datay_ids = datay_rank[rank_train_mask].long().view(-1, 1) - tstart_loss_calc = start_time(row_groups[0][0], rank) - - datay_ids = datay_rank[rank_train_mask].long() - - filtered_indices = torch.mul(datay_ids >= min_class, datay_ids < max_class).float() - indices = torch.nonzero(filtered_indices * torch.cuda.FloatTensor(datay_ids.size(), device=device).fill_(1)).squeeze() - # indices = torch.nonzero(filtered_indices * torch.FloatTensor(datay_ids.size(), device=device).fill_(1)).squeeze() - - datay_ids = datay_rank[rank_train_mask].long().view(-1, 1) - datay_ids = datay_ids.index_select(0, indices) - datay_ids -= min_class - outputs_ids = outputs.index_select(0, indices) - - # classes = torch.gather(outputs[rank_train_mask], 1, datay_ids) - classes = torch.gather(outputs_ids, 1, datay_ids) - loss_calc = torch.sum(classes) - loss_calc_tens = torch.Tensor([loss_calc.item()]) - - dist.all_reduce(loss_calc, op=dist.reduce_op.SUM, group=row_groups[rank_row][rank_c]) - - vertex_train_count = (data.train_mask.size(0) - (data.train_mask == 0).sum(dim=0)) - loss_calc = -loss_calc / vertex_train_count - - loss_calc_time += stop_time(row_groups[0][0], rank, tstart_loss_calc) - - loss_calc.backward() - - del filtered_indices - del indices - del datay_ids - del outputs_ids - del classes - del loss_calc - # loss = F.nll_loss(outputs[rank_train_mask], datay_rank[rank_train_mask]) - # loss.backward() - # print("loss: " + str(loss), flush=True) - else: - fake_loss = (outputs * torch.cuda.FloatTensor(outputs.size(), device=device).fill_(0)).sum() - # fake_loss = (outputs * torch.FloatTensor(outputs.size(), device=device).fill_(0)).sum() - # fake_loss = (outputs * torch.zeros(outputs.size())).sum() - fake_loss.backward() - - optimizer.step() - - return outputs - # del outputs - -def test(outputs, data, vertex_count, rank): - logits, accs = outputs, [] - datay_rank = torch.split(data.y, vertex_count)[rank] - for _, mask in data('train_mask', 'val_mask', 'test_mask'): - mask_rank = torch.split(mask, vertex_count)[rank] - count = mask_rank.nonzero().size(0) - if count > 0: - pred = logits[mask_rank].max(1)[1] - acc = pred.eq(datay_rank[mask_rank]).sum().item() / mask_rank.sum().item() - # pred = logits[mask].max(1)[1] - # acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item() - else: - acc = -1 - accs.append(acc) - return accs - - -# Split a COO into partitions of size n_per_proc -# Basically torch.split but for Sparse Tensors since pytorch doesn't support that. -def split_coo(adj_matrix, node_count, n_per_proc, dim, size): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - vtx_indices = list(range(0, node_count, n_per_proc)) - vtx_indices = vtx_indices[:proc_c] - vtx_indices.append(node_count) - - am_partitions = [] - for i in range(len(vtx_indices) - 1): - am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] - am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] - am_part[dim] -= vtx_indices[i] - am_partitions.append(am_part) - - return am_partitions, vtx_indices - # proc_row = proc_row_size(size) - # proc_col = proc_col_size(size) - - # vtx_indices = list(range(0, node_count, n_per_proc)) - # vtx_indices = vtx_indices[:proc_row] - # vtx_indices.append(node_count) - - # am_partitions = [] - # for i in range(len(vtx_indices) - 1): - # am_part = adj_matrix[:,(adj_matrix[dim,:] >= vtx_indices[i]).nonzero().squeeze(1)] - # am_part = am_part[:,(am_part[dim,:] < vtx_indices[i + 1]).nonzero().squeeze(1)] - # am_part[dim] -= vtx_indices[i] - # am_partitions.append(am_part) - - # return am_partitions, vtx_indices - -# Normalize all elements according to KW's normalization rule -def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx): - if not normalization: - return - - # Scale each edge (u, v) by 1 / (sqrt(u) * sqrt(v)) - indices = adj_part._indices() - values = adj_part._values() - - deg_map = dict() - for i in range(adj_part._nnz()): - u = indices[0][i] + row_vtx - v = indices[1][i] + col_vtx - - if u.item() in deg_map: - degu = deg_map[u.item()] - else: - degu = (adj_matrix[0] == u).sum().item() - deg_map[u.item()] = degu - - if v.item() in deg_map: - degv = deg_map[v.item()] - else: - degv = (adj_matrix[0] == v).sum().item() - deg_map[v.item()] = degv - - values[i] = values[i] / (math.sqrt(degu) * math.sqrt(degv)) - - # deg = torch.histc(adj_matrix[0].float(), bins=node_count) - # deg = deg.pow(-0.5) - - # row_len = adj_part.size(0) - # col_len = adj_part.size(1) - - # dleft = torch.sparse_coo_tensor([np.arange(row_vtx, row_vtx + row_len).tolist(), - # np.arange(row_vtx, row_vtx + row_len).tolist()], - # deg[row_vtx:(row_vtx + row_len)], - # size=(row_len, row_len), - # requires_grad=False) - - # dright = torch.sparse_coo_tensor([np.arange(col_vtx, col_vtx + col_len).tolist(), - # np.arange(col_vtx, col_vtx + col_len).tolist()], - # deg[row_vtx:(col_vtx + col_len)], - # size=(col_len, col_len), - # requires_grad=False) - - # adj_part = torch.sparse.mm(torch.sparse.mm(dleft, adj_part), dright) - # return adj_part - -def proc_row_size(size): - cube_root = int(size ** (1./ 3.)) - if cube_root ** 3 == size: - return cube_root - elif (cube_root + 1) ** 3 == size: - return cube_root + 1 - else: - print(f"CUBE ROOT ERROR") - -def proc_col_size(size): - cube_root = int(size ** (1./ 3.)) - if cube_root ** 3 == size: - return cube_root - elif (cube_root + 1) ** 3 == size: - return cube_root + 1 - else: - print(f"CUBE ROOT ERROR") - -def proc_c_size(size): - cube_root = int(size ** (1./ 3.)) - if cube_root ** 3 == size: - return cube_root - elif (cube_root + 1) ** 3 == size: - return cube_root + 1 - else: - print(f"CUBE ROOT ERROR") - -def twod_partition(rank, size, inputs, adj_matrix, data, features, classes, device): - node_count = inputs.size(0) - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - # n_per_proc = math.ceil(float(node_count) / proc_row) - n_per_proc = node_count // proc_row - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - - am_partitions = None - am_pbyp = None - - # Compute the adj_matrix and inputs partitions for this process - # TODO: Maybe I do want grad here. Unsure. - with torch.no_grad(): - # Column partitions - am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1, size) - - proc_node_count = vtx_indices[rank_col + 1] - vtx_indices[rank_col] - am_pbyp, _ = split_coo(am_partitions[rank_col], node_count, n_per_proc, 0, size) - for i in range(len(am_pbyp)): - if i == proc_row - 1: - last_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(last_node_count, proc_node_count), - requires_grad=False) - - scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank_col]) - else: - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(n_per_proc, proc_node_count), - requires_grad=False) - - scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], - vtx_indices[rank_col]) - - # input_rowparts = torch.split(inputs, math.ceil(float(inputs.size(0)) / proc_row), dim=0) - inputs_per_row = inputs.size(0) // proc_row - inputs_per_col = inputs.size(1) // proc_col - chunks_per_row = [] - chunks_per_col = [] - for i in range(proc_row): - if i == proc_row - 1: - chunks_per_row.append(inputs.size(0) - inputs_per_row * (proc_row - 1)) - else: - chunks_per_row.append(inputs_per_row) - for i in range(proc_col): - if i == proc_col - 1: - chunks_per_col.append(inputs.size(1) - inputs_per_col * (proc_col - 1)) - else: - chunks_per_col.append(inputs_per_col) - - # input_rowparts = torch.split(inputs, math.ceil(float(inputs.size(0)) / proc_row), dim=0) - input_rowparts = torch.split(inputs, chunks_per_row, dim=0) - input_partitions = [] - for i in input_rowparts: - # input_partitions.append(torch.split(i, math.ceil(float(inputs.size(1)) / proc_col), - # dim=1)) - input_partitions.append(torch.split(i, chunks_per_col, dim=1)) - - adj_matrix_loc = am_pbyp[rank_row] - inputs_loc = input_partitions[rank_row][rank_col] - - return inputs_loc, adj_matrix_loc, am_pbyp - -def threed_partition_loc(rank, size, inputs, adj_matrix, height, width, data, features, classes, device): - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - n_per_proc = width // proc_c - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - am_partitions, vtx_indices = split_coo(adj_matrix, width, n_per_proc, 1, size) - - for i in range(len(am_partitions)): - if i == proc_c - 1: - last_node_count = vtx_indices[i + 1] - vtx_indices[i] - am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], torch.ones(am_partitions[i].size(1)), - size=(height, last_node_count), - requires_grad=False) - - else: - am_partitions[i] = torch.sparse_coo_tensor(am_partitions[i], torch.ones(am_partitions[i].size(1)), - size=(height, n_per_proc), - requires_grad=False) - - inputs_per_row = inputs.size(0) // proc_c - chunks_per_row = [] - for i in range(proc_c): - if i == proc_row - 1: - chunks_per_row.append(inputs.size(0) - inputs_per_row * (proc_c - 1)) - else: - chunks_per_row.append(inputs_per_row) - - input_rowparts = torch.split(inputs, chunks_per_row, dim=0) - - adj_matrix_loc = am_partitions[rank_c] - inputs_loc = input_rowparts[rank_c] - - print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) - print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}", flush=True) - - return inputs_loc, adj_matrix_loc - -def threed_partition(rank, size, inputs, adj_matrix, data, features, classes, device): - node_count = inputs.size(0) - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - n_per_proc_row = node_count // proc_row - # n_per_proc_col = node_count // (proc_col * proc_c) - n_per_proc_col = node_count // proc_col - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - am_partitions = None - am_pbyp = None - - # Compute the adj_matrix and inputs partitions for this process - # TODO: Maybe I do want grad here. Unsure. - with torch.no_grad(): - # Compute column partitions - am_partitions, vtx_indices_col = split_coo(adj_matrix, node_count, n_per_proc_col, 1, size) - - proc_node_count = vtx_indices_col[rank_col * proc_c + rank_c + 1] - \ - vtx_indices_col[rank_col * proc_c + rank_c] - - # Compute row partitions - am_pbyp, vtx_indices_row = split_coo(am_partitions[rank_col * proc_c + rank_c], node_count, - n_per_proc_row, 0, size) - - for i in range(len(am_pbyp)): - if i == proc_row - 1: - last_node_count = vtx_indices_row[i + 1] - vtx_indices_row[i] - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(last_node_count, proc_node_count), - requires_grad=False) - - scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices_row[i], - vtx_indices_col[rank_col * proc_c + rank_c]) - else: - am_pbyp[i] = torch.sparse_coo_tensor(am_pbyp[i], torch.ones(am_pbyp[i].size(1)), - size=(n_per_proc_row, proc_node_count), - requires_grad=False) - - scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices_row[i], - vtx_indices_col[rank_col * proc_c + rank_c]) - - # input_rowparts = torch.split(inputs, math.ceil(float(inputs.size(0)) / proc_row), dim=0) - inputs_per_row = inputs.size(0) // (proc_row * proc_c) - inputs_per_col = inputs.size(1) // proc_col - chunks_per_row = [] - chunks_per_col = [] - for i in range(proc_row * proc_c): - if i == proc_row - 1: - chunks_per_row.append(inputs.size(0) - inputs_per_row * (proc_row * proc_c - 1)) - else: - chunks_per_row.append(inputs_per_row) - for i in range(proc_col): - if i == proc_col - 1: - chunks_per_col.append(inputs.size(1) - inputs_per_col * (proc_col - 1)) - else: - chunks_per_col.append(inputs_per_col) - - # input_rowparts = torch.split(inputs, math.ceil(float(inputs.size(0)) / proc_row), dim=0) - input_rowparts = torch.split(inputs, chunks_per_row, dim=0) - input_partitions = [] - for i in input_rowparts: - # input_partitions.append(torch.split(i, math.ceil(float(inputs.size(1)) / proc_col), - # dim=1)) - input_partitions.append(torch.split(i, chunks_per_col, dim=1)) - - adj_matrix_loc = am_pbyp[rank_row] - inputs_loc = input_partitions[rank_row * proc_c + rank_c][rank_col] - - print(f"rank: {rank} adj_matrix_loc.size: {adj_matrix_loc.size()}", flush=True) - print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}", flush=True) - return inputs_loc, adj_matrix_loc, am_pbyp - -def rank_to_devid(rank, acc_per_rank): - return rank % acc_per_rank - -def run(rank, size, inputs, adj_matrix, data, features, mid_layer, classes, device, acc_per_rank): - global comm_time - global comp_time - global epochs - global timing - - best_val_acc = test_acc = 0 - outputs = None - - group = dist.new_group(list(range(size))) - row_groups, col_groups, c_groups = get_proc_groups(rank, size, group) - - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) - proc_c = proc_c_size(size) - - rank_row = int((rank // proc_c) // proc_col) # i in process grid - rank_col = int((rank // proc_c) % proc_col) # j in process grid - rank_c = rank - (rank_row * (proc_col * proc_c) + rank_col * proc_c) # k in process grid - - if rank_row >= proc_row or rank_col >= proc_col or rank_c >= proc_c: - return - - rank_t = rank_col * proc_col * proc_c + rank_row * proc_c + rank_c - - transpose_groups = [] - transpose_group = None - - for i in range(proc_row): - transpose_groups_row = [] - for j in range(proc_col): - transpose_groups_col = [] - for k in range(proc_c): - local_rank = i * proc_col * proc_c + j * proc_c + k - local_rank_t = j * proc_col * proc_c + i * proc_c + k - - if local_rank < local_rank_t: - transpose_groups_col.append(dist.new_group([local_rank, local_rank_t])) - else: - transpose_groups_col.append(None) - transpose_groups_row.append(transpose_groups_col) - transpose_groups.append(transpose_groups_row) - - if rank < rank_t: - transpose_group = transpose_groups[rank_row][rank_col][rank_c] - else: - transpose_group = transpose_groups[rank_col][rank_row][rank_c] - - # adj_matrix_loc = torch.rand(node_count, n_per_proc) - # inputs_loc = torch.rand(n_per_proc, inputs.size(1)) - - torch.manual_seed(0) - weight1_nonleaf = torch.rand(features, mid_layer, requires_grad=True) - weight1_nonleaf = weight1_nonleaf.to(device) - weight1_nonleaf.retain_grad() - - weight2_nonleaf = torch.rand(mid_layer, classes, requires_grad=True) - weight2_nonleaf = weight2_nonleaf.to(device) - weight2_nonleaf.retain_grad() - - weight1 = Parameter(weight1_nonleaf) - weight2 = Parameter(weight2_nonleaf) - - optimizer = torch.optim.Adam([weight1, weight2], lr=0.01) - - # inputs_loc, adj_matrix_loc, _ = threed_partition(rank, size, inputs, adj_matrix, data, features, - # classes, device) - print(f"Before partitioning...", flush=True) - inputs_loc, adj_matrix_loc, _ = twod_partition(rank, size, inputs, adj_matrix, data, features, - classes, device) - adj_matrix_loc = adj_matrix_loc.coalesce() - - inputs_loc, adj_matrix_loc = threed_partition_loc(rank, size, inputs_loc, adj_matrix_loc.indices(), - adj_matrix_loc.size(0), adj_matrix_loc.size(1), - data, features, classes, device) - print(f"After partitioning...", flush=True) - - adj_matrix_loc = adj_matrix_loc.coalesce() - - inputs_loc = inputs_loc.to(device) - adj_matrix_loc = adj_matrix_loc.to(device) - - print(f"rank: {rank} Before first epoch...", flush=True) - # Do not time first epoch - timing_on = timing == True - timing = False - outputs = train(inputs_loc, weight1, weight2, inputs.size(0), adj_matrix_loc, None, - optimizer, data, rank, size, acc_per_rank, group, row_groups, - col_groups, transpose_group, c_groups) - print(f"After first epoch...", flush=True) - - if timing_on: - timing = True - - dist.barrier(group) - if rank == 0: - tstart = time.time() - - print(f"rank: {rank} Starting training...", flush=True) - for epoch in range(1, epochs): - if rank == 0: - tstart_epoch = time.time() - - outputs = train(inputs_loc, weight1, weight2, inputs.size(0), adj_matrix_loc, None, - optimizer, data, rank, size, acc_per_rank, group, row_groups, - col_groups, transpose_group, c_groups) - - # sync_and_sleep(rank, device) - if rank == 0: - # tstop_epoch = time.time() - # tstop = time.time() - # print(f"Epoch time: {epoch} {tstop_epoch - tstart_epoch}", flush=True) - # print("Total Time: " + str(tstop - tstart)) - # print("comm_time: " + str(comm_time)) - # print("comp_time: " + str(comp_time)) - # print(f"summa_sparse_comp: {summa_sparse_comp}") - # print(f"summa_sparse_bcast1: {summa_sparse_bcast1}") - # print(f"summa_sparse_bcast1_words: {summa_sparse_bcast1_words}") - # print(f"summa_sparse_bcast2: {summa_sparse_bcast2}") - # print(f"summa_sparse_reduce: {summa_sparse_reduce}") - # print(f"summa_sparse_bcast2_fwd: {summa_sparse_bcast2_fwd}") - # print(f"summa_sparse_bcast2_bwd: {summa_sparse_bcast2_bwd}") - # print(f"summa_sparse_bcast2_words: {summa_sparse_bcast2_words}") - # print(f"summa_comp: {summa_comp}") - # print(f"summa_bcast1: {summa_bcast1}") - # print(f"summa_bcast2: {summa_bcast2}") - # print(f"summa_reduce: {summa_reduce}") - # print(f"summa_loc_bcast: {summa_loc_bcast}") - # print(f"fwd_time: {fwd_time}") - # print(f"bwd_time: {bwd_time}") - # print(f"transpose_time: {transpose_time}") - # print(f"grad_weight_time: {grad_weight_time}") - # print(f"loss_calc_time: {loss_calc_time}") - # print(f"summa_sparse_time: {summa_sparse_time}") - # print(f"summa_time: {summa_time}") - print("Epoch: {:03d}".format(epoch), flush=True) - - dist.barrier() - if rank == 0: - tstop = time.time() - print("Time: " + str(tstop - tstart)) - - if rank == 0: - print("comm_time: " + str(comm_time)) - print("comp_time: " + str(comp_time)) - print(f"summa_sparse_comp: {summa_sparse_comp}") - print(f"summa_sparse_bcast1: {summa_sparse_bcast1}") - print(f"summa_sparse_bcast1_words: {summa_sparse_bcast1_words}") - print(f"summa_sparse_bcast2: {summa_sparse_bcast2}") - print(f"summa_sparse_reduce: {summa_sparse_reduce}") - print(f"summa_sparse_bcast2_fwd: {summa_sparse_bcast2_fwd}") - print(f"summa_sparse_bcast2_bwd: {summa_sparse_bcast2_bwd}") - print(f"summa_sparse_bcast2_words: {summa_sparse_bcast2_words}") - print(f"summa_comp: {summa_comp}") - print(f"summa_bcast1: {summa_bcast1}") - print(f"summa_bcast2: {summa_bcast2}") - print(f"summa_reduce: {summa_reduce}") - print(f"summa_loc_bcast: {summa_loc_bcast}") - print(f"fwd_time: {fwd_time}") - print(f"bwd_time: {bwd_time}") - print(f"transpose_time: {transpose_time}") - print(f"grad_weight_time: {grad_weight_time}") - print(f"loss_calc_time: {loss_calc_time}") - print(f"summa_sparse_time: {summa_sparse_time}") - print(f"summa_time: {summa_time}") - print(f"summa_loc_time: {summa_loc_time}") - - # All-gather outputs to test accuracy - # output_parts = [] - # for i in range(size): - # output_parts.append(torch.cuda.FloatTensor(am_partitions[0].size(1), classes).fill_(0)) - - # dist.all_gather(output_parts, outputs) - # outputs = torch.cat(output_parts, dim=0) - - # train_acc, val_acc, tmp_test_acc = test(outputs, data, am_partitions[0].size(1), rank) - # if val_acc > best_val_acc: - # best_val_acc = val_acc - # test_acc = tmp_test_acc - # log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}' - - # print(log.format(200, train_acc, best_val_acc, test_acc)) - print(f"rank: {rank} {outputs}") - return outputs - -def init_process(rank, size, inputs, adj_matrix, data, features, mid_layer, classes, device, - outputs, acc_per_rank, fn): - - run_outputs = fn(rank, size, inputs, adj_matrix, data, features, mid_layer, classes, device, - acc_per_rank) - if outputs is not None: - outputs[rank] = run_outputs.detach() - -def main(): - # graphname = 'Reddit' - global graphname - global mid_layer - - path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) - - # mid_layer = 16 - if graphname == 'Cora': - dataset = Planetoid(path, graphname, T.NormalizeFeatures()) - data = dataset[0] - num_features = dataset.num_features - num_classes = dataset.num_classes - elif graphname == 'Reddit': - dataset = Reddit(path, T.NormalizeFeatures()) - data = dataset[0] - num_features = dataset.num_features - num_classes = dataset.num_classes + 9 - elif graphname == 'Amazon': - # edge_index = torch.load(path + "/processed/amazon_graph.pt") - # edge_index = torch.load("/gpfs/alpine/bif115/scratch/alokt/Amazon/processed/amazon_graph_random.pt") - # edge_index = torch.load("/gpfs/alpine/bif115/scratch/alokt/Amazon/processed/amazon_large_randomized.pt") - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/Amazon/processed/data.pt") - print(f"Done loading coo", flush=True) - # edge_index = edge_index.t_() - # n = 9430086 - n = 14249639 - num_features = 300 - num_classes = 25 - # mid_layer = 24 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1) - data.train_mask = torch.ones(n).long() - elif graphname == 'subgraph5': - path = "/gpfs/alpine/bif115/scratch/alokt/HipMCL/" - edge_index = torch.load(path + "/processed/subgraph5_graph.pt") - n = 2186385 - num_features = 128 - # mid_layer = 64 - num_classes = 256 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1) - data.train_mask = torch.ones(n).long() - elif graphname == 'subgraph3': - # path = "/gpfs/alpine/bif115/scratch/alokt/HipMCL/" - # edge_index = torch.load(path + "/processed/subgraph3_graph.pt") - print(f"Loading coo...", flush=True) - edge_index = torch.load("../data/subgraph3/processed/data.pt") - print(f"Done loading coo", flush=True) - n = 8745542 - num_features = 128 - # mid_layer = 512 - # mid_layer = 64 - num_classes = 256 - inputs = torch.rand(n, num_features) - data = Data() - data.y = torch.rand(n).uniform_(0, num_classes - 1) - data.train_mask = torch.ones(n).long() - - if "OMPI_COMM_WORLD_RANK" in os.environ.keys(): - os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"] - dist.init_process_group(backend='nccl') - # dist.init_process_group('gloo', init_method='env://') - rank = dist.get_rank() - size = dist.get_world_size() - - mp.set_start_method('spawn', force=True) - # device = torch.device('cpu') - devid = rank_to_devid(rank, acc_per_rank) - - device = torch.device('cuda:{}'.format(devid)) - # device = torch.device('cpu') - print("device: " + str(device), flush=True) - torch.cuda.set_device(device) - curr_devid = torch.cuda.current_device() - devcount = torch.cuda.device_count() - print(f"devid: {curr_devid} {devcount}", flush=True) - - if graphname == "Amazon": - # edge_index = edge_index.to(device) - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - # inputs = inputs.to(device) - # inputs.requires_grad = True - data.y = data.y.to(device) - elif graphname == "subgraph5": - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - # inputs.requires_grad = True - data.y = data.y.to(device) - elif graphname == "subgraph3": - print(f"edge_index.size: {edge_index.size()}", flush=True) - data = data.to(device) - # inputs.requires_grad = True - data.y = data.y.to(device) - else: - data = data.to(device) - # data.x.requires_grad = True - inputs = data.x.to(device) - # inputs.requires_grad = True - data.y = data.y.to(device) - - edge_index = data.edge_index - print("edge count: " + str(len(edge_index[0]))) - - if normalization: - adj_matrix, _ = add_remaining_self_loops(edge_index) - else: - adj_matrix = edge_index - - outputs = None - print("Processes: " + str(size), flush=True) - - # init_process(rank, size, inputs, adj_matrix, data, dataset.num_features, mid_layer, dataset.num_classes, - init_process(rank, size, inputs, adj_matrix, data, num_features, mid_layer, num_classes, - device, outputs, acc_per_rank, run) - - if outputs is not None: - return outputs[0] - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument("--local_rank", type=int) - parser.add_argument("--accperrank", type=int) - parser.add_argument("--epochs", type=int) - parser.add_argument("--graphname", type=str) - parser.add_argument("--timing", type=str) - parser.add_argument("--midlayer", type=int) - args = parser.parse_args() - print(args) - - acc_per_rank = args.accperrank - if acc_per_rank is None: - acc_per_rank = 1 - - epochs = args.epochs - graphname = args.graphname - timing = args.timing == "True" - mid_layer = args.midlayer - - if (epochs is None) or (graphname is None) or (timing is None) or (mid_layer is None): - print(f"Error: missing argument {epochs} {graphname} {timing} {mid_layer}") - exit() - - print(f"Arguments: epochs: {epochs} graph: {graphname} timing: {timing} mid: {mid_layer}") - - print(main()) diff --git a/reddit.py b/reddit.py deleted file mode 100644 index f665eaf..0000000 --- a/reddit.py +++ /dev/null @@ -1,70 +0,0 @@ -import os -import os.path as osp - -import torch -import numpy as np -import scipy.sparse as sp -from torch_sparse import coalesce -from torch_geometric.data import (InMemoryDataset, Data, download_url, - extract_zip) - - -class Reddit(InMemoryDataset): - r"""The Reddit dataset from the `"Inductive Representation Learning on - Large Graphs" `_ paper, containing - Reddit posts belonging to different communities. - - Args: - root (string): Root directory where the dataset should be saved. - transform (callable, optional): A function/transform that takes in an - :obj:`torch_geometric.data.Data` object and returns a transformed - version. The data object will be transformed before every access. - (default: :obj:`None`) - pre_transform (callable, optional): A function/transform that takes in - an :obj:`torch_geometric.data.Data` object and returns a - transformed version. The data object will be transformed before - being saved to disk. (default: :obj:`None`) - """ - - url = 'https://data.dgl.ai/dataset/reddit.zip' - - def __init__(self, root, transform=None, pre_transform=None): - super(Reddit, self).__init__(root, transform, pre_transform) - self.data, self.slices = torch.load(self.processed_paths[0]) - - @property - def raw_file_names(self): - return ['reddit_data.npz', 'reddit_graph.npz'] - - @property - def processed_file_names(self): - return 'data.pt' - - def download(self): - path = download_url(self.url, self.raw_dir) - extract_zip(path, self.raw_dir) - os.unlink(path) - - def process(self): - data = np.load(osp.join(self.raw_dir, 'reddit_data.npz')) - x = torch.from_numpy(data['feature']).to(torch.float) - y = torch.from_numpy(data['label']).to(torch.long) - split = torch.from_numpy(data['node_types']) - - adj = sp.load_npz(osp.join(self.raw_dir, 'reddit_graph.npz')) - row = torch.from_numpy(adj.row).to(torch.long) - col = torch.from_numpy(adj.col).to(torch.long) - edge_index = torch.stack([row, col], dim=0) - edge_index, _ = coalesce(edge_index, None, x.size(0), x.size(0)) - - data = Data(x=x, edge_index=edge_index, y=y) - data.train_mask = split == 1 - data.val_mask = split == 2 - data.test_mask = split == 3 - - data = data if self.pre_transform is None else self.pre_transform(data) - - torch.save(self.collate([data]), self.processed_paths[0]) - - def __repr__(self): - return '{}()'.format(self.__class__.__name__) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..78bb7ba --- /dev/null +++ b/requirements.txt @@ -0,0 +1,151 @@ +WARNING: The directory '/global/cscratch1/sd/alokt/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag. +absl-py==0.13.0 +apex==0.1 +argon2-cffi @ file:///tmp/build/80754af9/argon2-cffi_1613037097816/work +astunparse==1.6.3 +async-generator @ file:///home/ktietz/src/ci/async_generator_1611927993394/work +attrs @ file:///tmp/build/80754af9/attrs_1620827162558/work +backcall @ file:///home/ktietz/src/ci/backcall_1611930011877/work +bleach @ file:///tmp/build/80754af9/bleach_1628110601003/work +blessings==1.7 +Bottleneck==1.3.2 +brotlipy==0.7.0 +cachetools==4.2.2 +certifi==2021.5.30 +cffi @ file:///tmp/build/80754af9/cffi_1625807838443/work +charset-normalizer @ file:///tmp/build/80754af9/charset-normalizer_1630003229654/work +click==8.0.1 +configparser==5.0.2 +cryptography @ file:///tmp/build/80754af9/cryptography_1616769286105/work +cycler==0.10.0 +debugpy @ file:///tmp/build/80754af9/debugpy_1629222698064/work +decorator @ file:///tmp/build/80754af9/decorator_1621259047763/work +defusedxml @ file:///tmp/build/80754af9/defusedxml_1615228127516/work +docker-pycreds==0.4.0 +entrypoints==0.3 +filelock==3.0.12 +fonttools==4.25.0 +future==0.18.2 +gitdb==4.0.7 +GitPython==3.1.18 +google-auth==1.35.0 +google-auth-oauthlib==0.4.5 +googledrivedownloader==0.4 +gpustat==0.6.0 +grpcio==1.39.0 +h5py==2.10.0 +idna @ file:///tmp/build/80754af9/idna_1622654382723/work +importlib-metadata @ file:///tmp/build/80754af9/importlib-metadata_1617874469820/work +iniconfig==1.1.1 +ipykernel @ file:///tmp/build/80754af9/ipykernel_1629995100487/work/dist/ipykernel-6.2.0-py3-none-any.whl +ipympl @ file:///tmp/build/80754af9/ipympl_1616528805894/work +ipython @ file:///tmp/build/80754af9/ipython_1628243924020/work +ipython_genutils==0.2.0 +ipywidgets @ file:///tmp/build/80754af9/ipywidgets_1610481889018/work +isodate==0.6.0 +jedi @ file:///tmp/build/80754af9/jedi_1611333125342/work +Jinja2 @ file:///tmp/build/80754af9/jinja2_1624781299557/work +joblib @ file:///tmp/build/80754af9/joblib_1613502643832/work +jsonschema @ file:///tmp/build/80754af9/jsonschema_1601490578438/work +jupyter-client @ file:///tmp/build/80754af9/jupyter_client_1616770841739/work +jupyter-core @ file:///tmp/build/80754af9/jupyter_core_1612213311222/work +jupyterlab-pygments @ file:///tmp/build/80754af9/jupyterlab_pygments_1601490720602/work +jupyterlab-widgets @ file:///tmp/build/80754af9/jupyterlab_widgets_1609884341231/work +kiwisolver @ file:///tmp/build/80754af9/kiwisolver_1612282420641/work +Markdown==3.3.4 +MarkupSafe @ file:///tmp/build/80754af9/markupsafe_1621528148836/work +matplotlib @ file:///tmp/build/80754af9/matplotlib-suite_1628003481308/work +matplotlib-inline @ file:///tmp/build/80754af9/matplotlib-inline_1628242447089/work +mistune==0.8.4 +mkl-fft==1.3.0 +mkl-random @ file:///tmp/build/80754af9/mkl_random_1626186064646/work +mkl-service==2.4.0 +mock @ file:///tmp/build/80754af9/mock_1607622725907/work +mpi4py==3.0.3 +msgpack==1.0.2 +munkres==1.1.4 +nbclient @ file:///tmp/build/80754af9/nbclient_1614364831625/work +nbconvert @ file:///tmp/build/80754af9/nbconvert_1624479060632/work +nbformat @ file:///tmp/build/80754af9/nbformat_1617383369282/work +nersc-tensorboard-helper @ git+https://github.com/NERSC/nersc-tensorboard-helper.git@b60a2acfc59b5a543250f4468cd842ecd7acda1c +nest-asyncio @ file:///tmp/build/80754af9/nest-asyncio_1613680548246/work +networkx==2.6.2 +notebook @ file:///tmp/build/80754af9/notebook_1629205607169/work +numexpr @ file:///tmp/build/80754af9/numexpr_1618856167419/work +numpy @ file:///tmp/build/80754af9/numpy_and_numpy_base_1626271506491/work +nvidia-ml-py3==7.352.0 +oauthlib==3.1.1 +olefile==0.46 +packaging @ file:///tmp/build/80754af9/packaging_1625611678980/work +pandas==1.3.2 +pandocfilters @ file:///tmp/build/80754af9/pandocfilters_1605120460739/work +parso @ file:///tmp/build/80754af9/parso_1617223946239/work +pathtools==0.1.2 +pexpect==4.8.0 +pickleshare==0.7.5 +Pillow @ file:///tmp/build/80754af9/pillow_1625670623564/work +pluggy==1.0.0 +prometheus-client @ file:///tmp/build/80754af9/prometheus_client_1623189609245/work +promise==2.3 +prompt-toolkit @ file:///tmp/build/80754af9/prompt-toolkit_1616415428029/work +protobuf==3.17.3 +psutil==5.8.0 +ptyprocess @ file:///tmp/build/80754af9/ptyprocess_1609355006118/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl +py==1.10.0 +pyasn1==0.4.8 +pyasn1-modules==0.2.8 +pycparser @ file:///tmp/build/80754af9/pycparser_1594388511720/work +Pygments @ file:///tmp/build/80754af9/pygments_1629234116488/work +pyOpenSSL @ file:///tmp/build/80754af9/pyopenssl_1608057966937/work +pyparsing @ file:///home/linux1/recipes/ci/pyparsing_1610983426697/work +pyrsistent @ file:///tmp/build/80754af9/pyrsistent_1600141720057/work +PySocks==1.7.1 +pytest==6.2.5 +python-dateutil @ file:///tmp/build/80754af9/python-dateutil_1626374649649/work +python-louvain==0.15 +pytz @ file:///tmp/build/80754af9/pytz_1612215392582/work +PyYAML==5.4.1 +pyzmq @ file:///tmp/build/80754af9/pyzmq_1628276010766/work +ray==1.6.0 +rdflib==6.0.0 +redis==3.5.3 +requests @ file:///tmp/build/80754af9/requests_1629994808627/work +requests-oauthlib==1.3.0 +rsa==4.7.2 +ruamel.yaml==0.17.16 +ruamel.yaml.clib==0.2.6 +scikit-learn @ file:///tmp/build/80754af9/scikit-learn_1621370412049/work +scipy @ file:///tmp/build/80754af9/scipy_1618855647378/work +Send2Trash==1.5.0 +sentry-sdk==1.3.1 +shortuuid==1.0.1 +sip==4.19.13 +six @ file:///tmp/build/80754af9/six_1623709665295/work +smmap==4.0.0 +sparse-coo-tensor-cpp==0.0.0 +subprocess32==3.5.4 +tables==3.6.1 +tensorboard==2.6.0 +tensorboard-data-server==0.6.1 +tensorboard-plugin-wit==1.8.0 +terminado==0.9.4 +testpath @ file:///tmp/build/80754af9/testpath_1624638946665/work +threadpoolctl @ file:///Users/ktietz/demo/mc3/conda-bld/threadpoolctl_1628668762525/work +toml==0.10.2 +torch==1.9.0 +torch-cluster==1.5.9 +torch-geometric==1.7.2 +torch-scatter==2.0.8 +torch-sparse==0.6.11 +torchvision==0.10.0a0+300a8a4 +tornado @ file:///tmp/build/80754af9/tornado_1606942300299/work +tqdm==4.62.2 +traitlets @ file:///home/ktietz/src/ci/traitlets_1611929699868/work +typing-extensions @ file:///tmp/build/80754af9/typing_extensions_1624965014186/work +urllib3 @ file:///tmp/build/80754af9/urllib3_1625084269274/work +wandb==0.12.1 +wcwidth @ file:///tmp/build/80754af9/wcwidth_1593447189090/work +webencodings==0.5.1 +Werkzeug==2.0.1 +widgetsnbextension==3.5.1 +zipp @ file:///tmp/build/80754af9/zipp_1625570634446/work diff --git a/shadow.py b/shadow.py new file mode 100644 index 0000000..daa1231 --- /dev/null +++ b/shadow.py @@ -0,0 +1,107 @@ +import copy +from typing import Optional + +import torch +from torch import Tensor + +from torch_geometric.data import Batch, Data +from torch_geometric.typing import WITH_TORCH_SPARSE, SparseTensor + + +class ShaDowKHopSampler(torch.utils.data.DataLoader): + r"""The ShaDow :math:`k`-hop sampler from the `"Decoupling the Depth and + Scope of Graph Neural Networks" `_ paper. + Given a graph in a :obj:`data` object, the sampler will create shallow, + localized subgraphs. + A deep GNN on this local graph then smooths the informative local signals. + + .. note:: + + For an example of using :class:`ShaDowKHopSampler`, see + `examples/shadow.py `_. + + Args: + data (torch_geometric.data.Data): The graph data object. + depth (int): The depth/number of hops of the localized subgraph. + num_neighbors (int): The number of neighbors to sample for each node in + each hop. + node_idx (LongTensor or BoolTensor, optional): The nodes that should be + considered for creating mini-batches. + If set to :obj:`None`, all nodes will be + considered. + replace (bool, optional): If set to :obj:`True`, will sample neighbors + with replacement. (default: :obj:`False`) + **kwargs (optional): Additional arguments of + :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size` or + :obj:`num_workers`. + """ + def __init__(self, data: Data, depth: int, num_neighbors: int, + node_idx: Optional[Tensor] = None, replace: bool = False, + **kwargs): + + if not WITH_TORCH_SPARSE: + raise ImportError( + f"'{self.__class__.__name__}' requires 'torch-sparse'") + + self.data = copy.copy(data) + self.depth = depth + self.num_neighbors = num_neighbors + self.replace = replace + + if data.edge_index is not None: + self.is_sparse_tensor = False + row, col = data.edge_index.cpu().share_memory_() + self.adj_t = SparseTensor( + row=row, col=col, value=torch.arange(col.size(0)), + sparse_sizes=(data.num_nodes, data.num_nodes)).t().share_memory_() + else: + self.is_sparse_tensor = True + self.adj_t = data.adj_t.cpu() + + if node_idx is None: + node_idx = torch.arange(self.adj_t.sparse_size(0)).share_memory_() + elif node_idx.dtype == torch.bool: + node_idx = node_idx.nonzero(as_tuple=False).view(-1).share_memory_() + self.node_idx = node_idx + + super().__init__(node_idx.tolist(), collate_fn=self.__collate__, + **kwargs) + + def __collate__(self, n_id): + print(f"n_id: {n_id}", flush=True) + n_id = torch.tensor(n_id) + + rowptr, col, value = self.adj_t.csr() + out = torch.ops.torch_sparse.ego_k_hop_sample_adj( + rowptr, col, n_id, self.depth, self.num_neighbors, self.replace) + rowptr, col, n_id, e_id, ptr, root_n_id = out + + adj_t = SparseTensor(rowptr=rowptr, col=col, + value=value[e_id] if value is not None else None, + sparse_sizes=(n_id.numel(), n_id.numel()), + is_sorted=True, trust_data=True) + + batch = Batch(batch=torch.ops.torch_sparse.ptr2ind(ptr, n_id.numel()), + ptr=ptr) + batch.root_n_id = root_n_id + + if self.is_sparse_tensor: + batch.adj_t = adj_t + else: + row, col, e_id = adj_t.t().coo() + batch.edge_index = torch.stack([row, col], dim=0) + + for k, v in self.data: + if k in ['edge_index', 'adj_t', 'num_nodes', 'batch', 'ptr']: + continue + if k == 'y' and v.size(0) == self.data.num_nodes: + batch[k] = v[n_id][root_n_id] + elif isinstance(v, Tensor) and v.size(0) == self.data.num_nodes: + batch[k] = v[n_id] + elif isinstance(v, Tensor) and v.size(0) == self.data.num_edges: + batch[k] = v[e_id] + else: + batch[k] = v + + return batch diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index 1fd5f96..d335145 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -2,6 +2,6 @@ from torch.utils import cpp_extension setup(name='sparse_coo_tensor_cpp', - ext_modules=[cpp_extension.CppExtension('sparse_coo_tensor_cpp', ['sparse_coo_tensor.cu'], - extra_compile_args=["-lcusparse"])], + ext_modules=[cpp_extension.CUDAExtension('sparse_coo_tensor_cpp', ['sparse_coo_tensor.cu'], + extra_compile_args=["-O2", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.7/../../math_libs/include/", "-I/global/u1/a/alokt/nsparse/cuda-c/inc", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/22.7/cuda/11.7/../../../22.7/examples/OpenACC/SDK/include/", "-lm", "-arch=sm_80", "-lcusparse"])], cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/sparse-extension/sparse_coo_tensor.cpp b/sparse-extension/sparse_coo_tensor.cpp deleted file mode 100644 index 9c906ef..0000000 --- a/sparse-extension/sparse_coo_tensor.cpp +++ /dev/null @@ -1,246 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include - -#include "cusparse.h" - -#include - -#include - -namespace py = pybind11; - -using namespace at::sparse; - -#define CHECK_CUSPARSE(func) \ -{ \ - cusparseStatus_t status = (func); \ - if (status != CUSPARSE_STATUS_SUCCESS) { \ - printf("CUSPARSE API failed at line %d with error: %s (%d)\n", \ - __LINE__, cusparseGetErrorString(status), status); \ - } \ -} - -#define CHECK_ERROR(str) \ - {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} - - -at::Tensor expand_values_if_needed(const at::Tensor& values) { - // expand - if (values.dim() == 0) { - // Mimic Numpy behavior here and treat it as a 1D tensor - return values.expand({1}); - } else { - return values; - } -} - -at::Tensor sparse_coo_tensor_gpu(const at::Tensor& indices, - const at::Tensor& values_, - at::ArrayRef size) { - - at::Tensor values = expand_values_if_needed(values_); - - int64_t sparse_dim = indices.size(0); - int64_t dense_dim = values.dim() - 1; - - return at::_sparse_coo_tensor_with_dims_and_tensors( - sparse_dim, dense_dim, size, indices, values, values.options().layout(at::kSparse)); -} - -template -void printCusparseDnMat(int64_t rows, int64_t cols, int64_t ld, T *values_dev) { - T* values_host = new T[rows*cols]; - cudaMemcpy(values_host, values_dev, rows*cols*sizeof(T), cudaMemcpyDeviceToHost); - for (int64_t row = 0; row < rows; row++) { - for (int64_t col = 0; col < cols; col++) { - // Cusparse dense matrices are stored in column-major order - std::cout << values_host[col*rows+row] << " "; - } - std::cout << std::endl; - } - std::cout << " values: "; - for (int64_t i = 0; i < rows*cols; i++) { - std::cout << values_host[i] << " "; - } - std::cout << std::endl; - std::cout << " shape: " << rows << ", " << cols << std::endl; - delete [] values_host; -} - -template -void printCusparseSpMat(int32_t rows, int32_t cols, int32_t nnz, int32_t *row_indices_dev, - int32_t *col_indices_dev, T *values_dev) { - T* values_host = new T[nnz]; - int32_t* row_indices_host = new int32_t[nnz]; - int32_t* col_indices_host = new int32_t[nnz]; - cudaMemcpy(values_host, values_dev, nnz*sizeof(T), cudaMemcpyDeviceToHost); - cudaMemcpy(row_indices_host, row_indices_dev, nnz*sizeof(int32_t), cudaMemcpyDeviceToHost); - cudaMemcpy(col_indices_host, col_indices_dev, nnz*sizeof(int32_t), cudaMemcpyDeviceToHost); - - for (int64_t i = 0; i < nnz; i++) { - std::cout << "(" << row_indices_host[i] - << ", " << col_indices_host[i] - << "): " << values_host[i] << std::endl; - } - std::cout << " values: "; - for (int64_t i = 0; i < nnz; i++) { - std::cout << values_host[i] << " "; - } - std::cout << std::endl; - std::cout << " row_indices: "; - for (int64_t i = 0; i < nnz; i++) { - std::cout << row_indices_host[i] << " "; - } - std::cout << std::endl; - std::cout << " col_indices: "; - for (int64_t i = 0; i < nnz; i++) { - std::cout << col_indices_host[i] << " "; - } - std::cout << std::endl; - std::cout << " shape: " << rows << ", " << cols << std::endl; - delete [] values_host; - delete [] row_indices_host; - delete [] col_indices_host; -} - -// at::Tensor spmm_gpu(const at::Tensor& A_rowindices, -void spmm_gpu(const at::Tensor& A_rowindices, - const at::Tensor& A_colindices, - const at::Tensor& A_values, - int32_t n, - int32_t m, - at::Tensor& B, - at::Tensor& C) { - - // cusparseHandle_t handle; - // CHECK_CUSPARSE(cusparseCreate(&handle)); - auto state = at::globalContext().lazyInitCUDA(); - // auto handle = THCState_getCurrentSparseHandle(state); - auto handle = at::cuda::getCurrentCUDASparseHandle(); - - // Impl1 -- coo2csr + csrmm2 - int nnz = A_values.size(0); - - clock_t start, stop; - - int32_t *d_a_csrrows; - - // int devid_old = 0; - // cudaGetDevice(&devid_old); - // cudaSetDevice(devid); - - cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); - CHECK_CUSPARSE(cusparseXcoo2csr(handle, - A_rowindices.data(), - nnz, - n, - d_a_csrrows, - CUSPARSE_INDEX_BASE_ZERO)); - - int32_t b_row = B.size(0); - int32_t b_col = B.size(1); - int32_t c_row = C.size(0); - int32_t c_col = C.size(1); - - float alpha = 1; - float beta = 1; - cusparseSpMatDescr_t matA; - CHECK_CUSPARSE(cusparseCreateCsr(&matA, - n, // rows - // b_col, // cols - m, // cols - nnz, // nnz - d_a_csrrows, // csrRowOffsets - A_colindices.data(), // csrColInd - A_values.data(), // csrValues - CUSPARSE_INDEX_32I, // csrRowOffsetsType - CUSPARSE_INDEX_32I, // csrColIndType - CUSPARSE_INDEX_BASE_ZERO, // idxBase, - CUDA_R_32F)); // valueType - - // Row-major to column-major - B.t_(); - B.set_data(B.contiguous()); - B.set_data(B.view({b_row, b_col})); - - cusparseDnMatDescr_t matB; - CHECK_CUSPARSE(cusparseCreateDnMat(&matB, - // b_col, // rows - b_row, // rows - // b_row, // cols - b_col, // cols - // b_col, // ld - b_row, // ld - B.data(), // values - CUDA_R_32F, // valueType - CUSPARSE_ORDER_COL)); // order - - // Row-major to column-major - C.t_(); - C.set_data(C.contiguous()); - C.set_data(C.view({c_row, c_col})); - - cusparseDnMatDescr_t matC; - CHECK_CUSPARSE(cusparseCreateDnMat(&matC, - n, // rows - B.size(1), // cols - // n, // ld - n, // ld - C.data(), // values - CUDA_R_32F, // valueType - CUSPARSE_ORDER_COL)); // order - - size_t bufferSize; - CHECK_CUSPARSE(cusparseSpMM_bufferSize(handle, // handle, - CUSPARSE_OPERATION_NON_TRANSPOSE, // opA - // CUSPARSE_OPERATION_TRANSPOSE, // opB - CUSPARSE_OPERATION_NON_TRANSPOSE, // opB - &alpha, // alpha - matA, // matA - matB, // matB - &beta, // beta - matC, // matC - CUDA_R_32F, // computeType - CUSPARSE_CSRMM_ALG1, // alg - &bufferSize)); // bufferSize - - - void* d_buffer = NULL; - CHECK_ERROR(cudaMalloc(&d_buffer, bufferSize)); - - CHECK_CUSPARSE(cusparseSpMM(handle, // handle, - CUSPARSE_OPERATION_NON_TRANSPOSE, // opA - // CUSPARSE_OPERATION_TRANSPOSE, // opB - CUSPARSE_OPERATION_NON_TRANSPOSE, // opB - &alpha, // alpha - matA, // matA - matB, // matB - &beta, // beta - matC, // matC - CUDA_R_32F, // computeType - CUSPARSE_CSRMM_ALG1, // alg - d_buffer)); // buffer - - - CHECK_ERROR(cudaFree(d_a_csrrows)); - CHECK_ERROR(cudaFree(d_buffer)); - - // Column-major to row-major - C.set_data(C.view({c_col, c_row})); - C.t_(); - - // Column-major to row-major - B.set_data(B.view({b_col, b_row})); - B.t_(); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("sparse_coo_tensor_gpu", &sparse_coo_tensor_gpu, "Sparse Tensor GPU-only constructor"); - m.def("spmm_gpu", &spmm_gpu, "SpMM wrapper for cusparse"); -} diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 355919f..bd42d2d 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -5,14 +5,27 @@ #include #include #include -#include +#include +#include +#include +#include #include "cusparse.h" #include +typedef float real; +#include "nsparse_asm.h" + +// #include #include +#include +#include +#include +#include +#include + namespace py = pybind11; using namespace at::sparse; @@ -29,6 +42,1380 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} +// NSPARSE defs +#define div_round_up(a, b) ((a % b == 0)? a / b : a / b + 1) +#define WARP_BIT 5 +#define WARP 32 +#define MAX_LOCAL_THREAD_NUM 1024 +#define MAX_THREAD_BLOCK (MAX_LOCAL_THREAD_NUM / WARP) +#define BIN_NUM 7 +#define PWARP 4 +#define IMB_PWMIN 32 +#define B_PWMIN 16 +#define IMB_MIN 512 +#define B_MIN 256 +#define IMB_PW_SH_SIZE 4096 +#define B_PW_SH_SIZE 2048 +#define IMB_SH_SIZE 1024 +#define B_SH_SIZE 512 +// #define HASH_SCAL 107 // Set disjoint number to COMP_SH_SIZE +#define HASH_SCAL 4 // Set disjoint number to COMP_SH_SIZE + +/* Structure for SpGEMM */ +typedef struct { + cudaStream_t *stream; + int *bin_size; + int *bin_offset; + int *d_bin_size; + int *d_bin_offset; + int *d_row_nz; + int *d_row_perm; + int max_intprod; + int max_nz; + int *d_max; +} sfBIN; + +void init_bin(sfBIN *bin, int M) +{ + int i; + bin->stream = (cudaStream_t *)malloc(sizeof(cudaStream_t) * BIN_NUM); + for (i = 0; i < BIN_NUM; i++) { + cudaStreamCreate(&(bin->stream[i])); + } + + bin->bin_size = (int *)malloc(sizeof(int) * BIN_NUM); + bin->bin_offset = (int *)malloc(sizeof(int) * BIN_NUM); + cudaMalloc((void **)&(bin->d_row_perm), sizeof(int) * M); + cudaMalloc((void **)&(bin->d_row_nz), sizeof(int) * (M + 1)); + cudaMalloc((void **)&(bin->d_max), sizeof(int)); + cudaMalloc((void **)&(bin->d_bin_size), sizeof(int) * BIN_NUM); + cudaMalloc((void **)&(bin->d_bin_offset), sizeof(int) * BIN_NUM); + + // auto options = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA); + // torch::Tensor d_row_perm_tens = torch::empty({M}, torch::TensorOptions() + // .dtype(torch::kInt32) + // .device(torch::kCUDA)); + // torch::Tensor d_row_nz_tens = torch::zeros({M + 1}, options); + // torch::Tensor d_max_tens = torch::zeros({1}, options); + // torch::Tensor d_bin_size_tens = torch::zeros({BIN_NUM}, options); + // torch::Tensor d_bin_offset_tens = torch::zeros({BIN_NUM}, options); + + // bin->d_row_perm = d_row_perm_tens.data(); + // bin->d_row_nz = d_row_nz_tens.data(); + // bin->d_max = d_max_tens.data(); + // bin->d_bin_size = d_bin_size_tens.data(); + // bin->d_bin_offset = d_bin_offset_tens.data(); + + i = 0; + bin->max_intprod = 0; + bin->max_nz = 0; + +} + +void release_bin(sfBIN bin) +{ + int i; + cudaFree(bin.d_row_perm); + cudaFree(bin.d_row_nz); + cudaFree(bin.d_max); + cudaFree(bin.d_bin_size); + cudaFree(bin.d_bin_offset); + + free(bin.bin_size); + free(bin.bin_offset); + for (i = 0; i < BIN_NUM; i++) { + cudaStreamDestroy(bin.stream[i]); + } + free(bin.stream); +} + +__global__ void set_intprod_num(int *d_arpt, int *d_acol, + const int* __restrict__ d_brpt, + int *d_row_intprod, int *d_max_intprod, + int M) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= M) { + return; + } + int nz_per_row = 0; + int j; + for (j = d_arpt[i]; j < d_arpt[i + 1]; j++) { + nz_per_row += d_brpt[d_acol[j] + 1] - d_brpt[d_acol[j]]; + } + d_row_intprod[i] = nz_per_row; + atomicMax(d_max_intprod, nz_per_row); +} + +__global__ void set_bin(int *d_row_nz, int *d_bin_size, int *d_max, + int M, int min, int mmin) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= M) { + return; + } + int nz_per_row = d_row_nz[i]; + + atomicMax(d_max, nz_per_row); + + int j = 0; + for (j = 0; j < BIN_NUM - 2; j++) { + if (nz_per_row <= (min << j)) { + if (nz_per_row <= (mmin)) { + atomicAdd(d_bin_size + j, 1); + } + else { + atomicAdd(d_bin_size + j + 1, 1); + } + return; + } + } + atomicAdd(d_bin_size + BIN_NUM - 1, 1); +} + +__global__ void init_row_perm(int *d_permutation, int M) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + + if (i >= M) { + return; + } + + d_permutation[i] = i; +} + +__global__ void set_row_perm(int *d_bin_size, int *d_bin_offset, + int *d_max_row_nz, int *d_row_perm, + int M, int min, int mmin) +{ + int i = blockDim.x * blockIdx.x + threadIdx.x; + if (i >= M) { + return; + } + + int nz_per_row = d_max_row_nz[i]; + int dest; + + int j = 0; + for (j = 0; j < BIN_NUM - 2; j++) { + if (nz_per_row <= (min << j)) { + if (nz_per_row <= mmin) { + dest = atomicAdd(d_bin_size + j, 1); + d_row_perm[d_bin_offset[j] + dest] = i; + } + else { + dest = atomicAdd(d_bin_size + j + 1, 1); + d_row_perm[d_bin_offset[j + 1] + dest] = i; + } + return; + } + } + dest = atomicAdd(d_bin_size + BIN_NUM - 1, 1); + d_row_perm[d_bin_offset[BIN_NUM - 1] + dest] = i; + +} + +void set_max_bin(int *d_arpt, int *d_acol, int *d_brpt, sfBIN *bin, int M) +{ + int i; + int GS, BS; + + for (i = 0; i < BIN_NUM; i++) { + bin->bin_size[i] = 0; + bin->bin_offset[i] = 0; + } + + cudaMemcpy(bin->d_bin_size, bin->bin_size, sizeof(int) * BIN_NUM, cudaMemcpyHostToDevice); + cudaMemcpy(bin->d_max, &(bin->max_intprod), sizeof(int), cudaMemcpyHostToDevice); + + BS = 1024; + GS = div_round_up(M, BS); + set_intprod_num<<>>(d_arpt, d_acol, d_brpt, bin->d_row_nz, bin->d_max, M); + cudaMemcpy(&(bin->max_intprod), bin->d_max, sizeof(int), cudaMemcpyDeviceToHost); + + if (bin->max_intprod > IMB_PWMIN) { + set_bin<<>>(bin->d_row_nz, bin->d_bin_size, bin->d_max, M, IMB_MIN, IMB_PWMIN); + + cudaMemcpy(bin->bin_size, bin->d_bin_size, sizeof(int) * BIN_NUM, cudaMemcpyDeviceToHost); + cudaMemcpy(bin->d_bin_size, bin->bin_offset, sizeof(int) * BIN_NUM, cudaMemcpyHostToDevice); + + for (i = 0; i < BIN_NUM - 1; i++) { + bin->bin_offset[i + 1] = bin->bin_offset[i] + bin->bin_size[i]; + } + cudaMemcpy(bin->d_bin_offset, bin->bin_offset, sizeof(int) * BIN_NUM, cudaMemcpyHostToDevice); + + set_row_perm<<>>(bin->d_bin_size, bin->d_bin_offset, bin->d_row_nz, bin->d_row_perm, M, IMB_MIN, IMB_PWMIN); + } + else { + bin->bin_size[0] = M; + for (i = 1; i < BIN_NUM; i++) { + bin->bin_size[i] = 0; + } + bin->bin_offset[0] = 0; + for (i = 1; i < BIN_NUM; i++) { + bin->bin_offset[i] = M; + } + init_row_perm<<>>(bin->d_row_perm, M); + } +} + + +void set_min_bin(sfBIN *bin, int M) +{ + int i; + int GS, BS; + + for (i = 0; i < BIN_NUM; i++) { + bin->bin_size[i] = 0; + bin->bin_offset[i] = 0; + } + + cudaMemcpy(bin->d_bin_size, bin->bin_size, sizeof(int) * BIN_NUM, cudaMemcpyHostToDevice); + cudaMemcpy(bin->d_max, &(bin->max_nz), sizeof(int), cudaMemcpyHostToDevice); + + BS = 1024; + GS = div_round_up(M, BS); + set_bin<<>>(bin->d_row_nz, bin->d_bin_size, + bin->d_max, + M, B_MIN, B_PWMIN); + + cudaMemcpy(&(bin->max_nz), bin->d_max, sizeof(int), cudaMemcpyDeviceToHost); + if (bin->max_nz > B_PWMIN) { + cudaMemcpy(bin->bin_size, bin->d_bin_size, sizeof(int) * BIN_NUM, cudaMemcpyDeviceToHost); + cudaMemcpy(bin->d_bin_size, bin->bin_offset, sizeof(int) * BIN_NUM, cudaMemcpyHostToDevice); + + for (i = 0; i < BIN_NUM - 1; i++) { + bin->bin_offset[i + 1] = bin->bin_offset[i] + bin->bin_size[i]; + } + cudaMemcpy(bin->d_bin_offset, bin->bin_offset, sizeof(int) * BIN_NUM, cudaMemcpyHostToDevice); + + set_row_perm<<>>(bin->d_bin_size, bin->d_bin_offset, bin->d_row_nz, bin->d_row_perm, M, B_MIN, B_PWMIN); + } + + else { + bin->bin_size[0] = M; + for (i = 1; i < BIN_NUM; i++) { + bin->bin_size[i] = 0; + } + bin->bin_offset[0] = 0; + for (i = 1; i < BIN_NUM; i++) { + bin->bin_offset[i] = M; + } + BS = 1024; + GS = div_round_up(M, BS); + init_row_perm<<>>(bin->d_row_perm, M); + } +} + +__global__ void init_value(real *d_val, int nz) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= nz) { + return; + } + d_val[i] = 0; +} + +__global__ void init_check(int *d_check, int nz) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= nz) { + return; + } + d_check[i] = -1; +} + +__global__ void set_row_nz_bin_pwarp(const int *d_arpt, const int *d_acol, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const int *d_row_perm, + int *d_row_nz, + int bin_offset, int M) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + int rid = i / PWARP; + int tid = i % PWARP; + int local_rid = rid % (blockDim.x / PWARP); + + int j, k; + int soffset; + int acol, bcol, key, hash, adr, nz, old; + __shared__ int check[IMB_PW_SH_SIZE]; + + soffset = local_rid * IMB_PWMIN; + + for (j = tid; j < IMB_PWMIN; j += PWARP) { + check[soffset + j] = -1; + } + if (rid >= M) { + return; + } + + rid = d_row_perm[rid + bin_offset]; + nz = 0; + for (j = d_arpt[rid] + tid; j < d_arpt[rid + 1]; j += PWARP) { + acol = ld_gbl_int32(d_acol + j); + for (k = d_brpt[acol]; k < d_brpt[acol + 1]; k++) { + bcol = d_bcol[k]; + key = bcol; + hash = (bcol * HASH_SCAL) & (IMB_PWMIN - 1); + adr = soffset + hash; + while (1) { + if (check[adr] == key) { + break; + } + else if (check[adr] == -1) { + old = atomicCAS(check + adr, -1, key); + if (old == -1) { + nz++; + break; + } + } + else { + hash = (hash + 1) & (IMB_PWMIN - 1); + adr = soffset + hash; + } + } + } + } + + for (j = PWARP / 2; j >= 1; j /= 2) { + // nz += __shfl_xor(nz, j); + nz += __shfl_xor_sync(0xFFFFFFFF, nz, j); + } + + if (tid == 0) { + d_row_nz[rid] = nz; + } +} + + +template +__global__ void set_row_nz_bin_each(const int *d_arpt, const int *d_acol, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const int *d_row_perm, + int *d_row_nz, int bin_offset, int M) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + int rid = i / WARP; + int tid = i % WARP; + int wid = rid % (blockDim.x / WARP); + int j, k, l; + int bcol, key, hash, old; + int nz, adr; + int acol, ccol; + int soffset; + + soffset = wid * SH_ROW; + + __shared__ int check[IMB_SH_SIZE]; + for (j = tid; j < SH_ROW; j += WARP) { + check[soffset + j] = -1; + } + + if (rid >= M) { + return; + } + + acol = 0; + nz = 0; + rid = d_row_perm[rid + bin_offset]; + for (j = d_arpt[rid]; j < d_arpt[rid + 1]; j += WARP) { + if (j + tid < d_arpt[rid + 1]) acol = ld_gbl_int32(d_acol + j + tid); + for (l = 0; l < WARP && j + l < d_arpt[rid + 1]; l++) { + ccol = __shfl(acol, l); + for (k = d_brpt[ccol] + tid; k < d_brpt[ccol + 1]; k += WARP) { + bcol = d_bcol[k]; + key = bcol; + hash = (bcol * HASH_SCAL) & (SH_ROW - 1); + adr = soffset + hash; + while (1) { + if (check[adr] == key) { + break; + } + else if (check[adr] == -1) { + old = atomicCAS(check + adr, -1, key); + if (old == -1) { + nz++; + break; + } + } + else { + hash = (hash + 1) & (SH_ROW - 1); + adr = soffset + hash; + } + } + } + } + } + + for (j = WARP / 2; j >= 1; j /= 2) { + // nz += __shfl_xor(nz, j); + nz += __shfl_xor_sync(0xFFFFFFFF, nz, j); + } + + if (tid == 0) { + d_row_nz[rid] = nz; + } +} + +template +__global__ void set_row_nz_bin_each_tb(const int *d_arpt, const int *d_acol, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + int *d_row_perm, int *d_row_nz, + int bin_offset, int M) +{ + int rid = blockIdx.x; + int tid = threadIdx.x & (WARP - 1); + int wid = threadIdx.x / WARP; + int wnum = blockDim.x / WARP; + int j, k; + int bcol, key, hash, old; + int nz, adr; + int acol; + + __shared__ int check[SH_ROW]; + for (j = threadIdx.x; j < SH_ROW; j += blockDim.x) { + check[j] = -1; + } + + if (rid >= M) { + return; + } + + __syncthreads(); + + nz = 0; + rid = d_row_perm[rid + bin_offset]; + for (j = d_arpt[rid] + wid; j < d_arpt[rid + 1]; j += wnum) { + acol = ld_gbl_int32(d_acol + j); + for (k = d_brpt[acol] + tid; k < d_brpt[acol + 1]; k += WARP) { + bcol = d_bcol[k]; + key = bcol; + hash = (bcol * HASH_SCAL) & (SH_ROW - 1); + adr = hash; + while (1) { + if (check[adr] == key) { + break; + } + else if (check[adr] == -1) { + old = atomicCAS(check + adr, -1, key); + if (old == -1) { + nz++; + break; + } + } + else { + hash = (hash + 1) & (SH_ROW - 1); + adr = hash; + } + } + } + } + + for (j = WARP / 2; j >= 1; j /= 2) { + // nz += __shfl_xor(nz, j); + nz += __shfl_xor_sync(0xFFFFFFFF, nz, j); + } + + __syncthreads(); + if (threadIdx.x == 0) { + check[0] = 0; + } + __syncthreads(); + + if (tid == 0) { + atomicAdd(check, nz); + } + __syncthreads(); + + if (threadIdx.x == 0) { + d_row_nz[rid] = check[0]; + } +} + +template +__global__ void set_row_nz_bin_each_tb_large(const int *d_arpt, const int *d_acol, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + int *d_row_perm, int *d_row_nz, + int *d_fail_count, int *d_fail_perm, + int bin_offset, int M) +{ + int rid = blockIdx.x; + int tid = threadIdx.x & (WARP - 1); + int wid = threadIdx.x / WARP; + int wnum = blockDim.x / WARP; + int j, k; + int bcol, key, hash, old; + int adr; + int acol; + + __shared__ int check[SH_ROW]; + __shared__ int snz[1]; + for (j = threadIdx.x; j < SH_ROW; j += blockDim.x) { + check[j] = -1; + } + if (threadIdx.x == 0) { + snz[0] = 0; + } + + if (rid >= M) { + return; + } + + __syncthreads(); + + rid = d_row_perm[rid + bin_offset]; + int count = 0; + int border = SH_ROW >> 1; + for (j = d_arpt[rid] + wid; j < d_arpt[rid + 1]; j += wnum) { + acol = ld_gbl_int32(d_acol + j); + for (k = d_brpt[acol] + tid; k < d_brpt[acol + 1]; k += WARP) { + bcol = d_bcol[k]; + key = bcol; + hash = (bcol * HASH_SCAL) & (SH_ROW - 1); + adr = hash; + while (count < border && snz[0] < border) { + if (check[adr] == key) { + break; + } + else if (check[adr] == -1) { + old = atomicCAS(check + adr, -1, key); + if (old == -1) { + atomicAdd(snz, 1); + break; + } + } + else { + hash = (hash + 1) & (SH_ROW - 1); + adr = hash; + count++; + } + } + if (count >= border || snz[0] >= border) { + break; + } + } + if (count >= border || snz[0] >= border) { + break; + } + } + + __syncthreads(); + if (count >= border || snz[0] >= border) { + if (threadIdx.x == 0) { + int d = atomicAdd(d_fail_count, 1); + d_fail_perm[d] = rid; + } + } + else { + if (threadIdx.x == 0) { + d_row_nz[rid] = snz[0]; + } + } +} + +__global__ void set_row_nz_bin_each_gl(const int *d_arpt, const int *d_acol, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const int *d_row_perm, + int *d_row_nz, int *d_check, + int max_row_nz, int bin_offset, int M) +{ + int rid = blockIdx.x; + int tid = threadIdx.x & (WARP - 1); + int wid = threadIdx.x / WARP; + int wnum = blockDim.x / WARP; + int j, k; + int bcol, key, hash, old; + int nz, adr; + int acol; + int offset = rid * max_row_nz; + + __shared__ int snz[1]; + if (threadIdx.x == 0) { + snz[0] = 0; + } + __syncthreads(); + + if (rid >= M) { + return; + } + + nz = 0; + rid = d_row_perm[rid + bin_offset]; + for (j = d_arpt[rid] + wid; j < d_arpt[rid + 1]; j += wnum) { + acol = ld_gbl_int32(d_acol + j); + for (k = d_brpt[acol] + tid; k < d_brpt[acol + 1]; k += WARP) { + bcol = d_bcol[k]; + key = bcol; + hash = (bcol * HASH_SCAL) % max_row_nz; + adr = offset + hash; + while (1) { + if (d_check[adr] == key) { + break; + } + else if (d_check[adr] == -1) { + old = atomicCAS(d_check + adr, -1, key); + if (old == -1) { + nz++; + break; + } + } + else { + hash = (hash + 1) % max_row_nz; + adr = offset + hash; + } + } + } + } + + for (j = WARP / 2; j >= 1; j /= 2) { + // nz += __shfl_xor(nz, j); + nz += __shfl_xor_sync(0xFFFFFFFF, nz, j); + } + + if (tid == 0) { + atomicAdd(snz, nz); + } + __syncthreads(); + if (threadIdx.x == 0) { + d_row_nz[rid] = snz[0]; + } +} + +void set_row_nnz(int *d_arpt, int *d_acol, + int *d_brpt, int *d_bcol, + int *d_crpt, + sfBIN *bin, + int M, int *nnz); + + +__global__ void calculate_value_col_bin_pwarp(const int *d_arpt, + const int *d_acol, + const real *d_aval, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const real* __restrict__ d_bval, + int *d_crpt, + int *d_ccol, + real *d_cval, + const int *d_row_perm, + int *d_nz, + int bin_offset, + int bin_size) { + + int i = blockIdx.x * blockDim.x + threadIdx.x; + int rid = i / PWARP; + int tid = i % PWARP; + int local_rid = rid % (blockDim.x / PWARP); + int j; + __shared__ int shared_check[B_PW_SH_SIZE]; + __shared__ real shared_value[B_PW_SH_SIZE]; + + int soffset = local_rid * (B_PWMIN); + + for (j = tid; j < (B_PWMIN); j += PWARP) { + shared_check[soffset + j] = -1; + shared_value[soffset + j] = 0; + } + + if (rid >= bin_size) { + return; + } + rid = d_row_perm[rid + bin_offset]; + + if (tid == 0) { + d_nz[rid] = 0; + } + int k; + int acol, bcol, hash, key, adr; + int offset = d_crpt[rid]; + int old, index; + real aval, bval; + + for (j = d_arpt[rid] + tid; j < d_arpt[rid + 1]; j += PWARP) { + acol = ld_gbl_int32(d_acol + j); + aval = ld_gbl_real(d_aval + j); + for (k = d_brpt[acol]; k < d_brpt[acol + 1]; k++) { + bcol = d_bcol[k]; + bval = d_bval[k]; + + key = bcol; + hash = (bcol * HASH_SCAL) & ((B_PWMIN) - 1); + adr = soffset + hash; + while (1) { + if (shared_check[adr] == key) { + atomic_fadd(shared_value + adr, aval * bval); + break; + } + else if (shared_check[adr] == -1) { + old = atomicCAS(shared_check + adr, -1, key); + if (old == -1) { + atomic_fadd(shared_value + adr, aval * bval); + break; + } + } + else { + hash = (hash + 1) & ((B_PWMIN) - 1); + adr = soffset + hash; + } + } + } + } + + for (j = tid; j < (B_PWMIN); j += PWARP) { + if (shared_check[soffset + j] != -1) { + index = atomicAdd(d_nz + rid, 1); + shared_check[soffset + index] = shared_check[soffset + j]; + shared_value[soffset + index] = shared_value[soffset + j]; + } + } + int nz = d_nz[rid]; + // Sorting for shared data + int count, target; + for (j = tid; j < nz; j += PWARP) { + target = shared_check[soffset + j]; + count = 0; + for (k = 0; k < nz; k++) { + count += (unsigned int)(shared_check[soffset + k] - target) >> 31; + } + d_ccol[offset + count] = shared_check[soffset + j]; + d_cval[offset + count] = shared_value[soffset + j]; + } +} + + +template +__global__ void calculate_value_col_bin_each(const int *d_arpt, + const int *d_acol, + const real *d_aval, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const real* __restrict__ d_bval, + int *d_crpt, + int *d_ccol, + real *d_cval, + const int *d_row_perm, + int *d_nz, + int bin_offset, + int bin_size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + int rid = i / WARP; + int tid = i % WARP; + int wid = rid % (blockDim.x / WARP); + int j; + __shared__ int shared_check[B_SH_SIZE]; + __shared__ real shared_value[B_SH_SIZE]; + + int soffset = wid * SH_ROW; + + for (j = tid; j < SH_ROW; j += WARP) { + shared_check[soffset + j] = -1; + shared_value[soffset + j] = 0; + } + + if (rid >= bin_size) { + return; + } + rid = d_row_perm[rid + bin_offset]; + + if (tid == 0) { + d_nz[rid] = 0; + } + int lacol, acol; + int k, l; + int bcol, hash, key, adr; + int offset = d_crpt[rid]; + int old, index; + real laval, aval, bval; + + lacol = 0; + for (j = d_arpt[rid]; j < d_arpt[rid + 1]; j += WARP) { + if (j + tid < d_arpt[rid + 1]) { + lacol = ld_gbl_int32(d_acol + j + tid); + laval = ld_gbl_real(d_aval + j + tid); + } + for (l = 0; l < WARP && j + l < d_arpt[rid + 1]; l++) { + acol = __shfl(lacol, l); + aval = __shfl(laval, l); + for (k = d_brpt[acol] + tid; k < d_brpt[acol + 1]; k += WARP) { + bcol = d_bcol[k]; + bval = d_bval[k]; + + key = bcol; + hash = (bcol * HASH_SCAL) & (SH_ROW - 1); + adr = soffset + hash; + while (1) { + if (shared_check[adr] == key) { + atomic_fadd(shared_value + adr, aval * bval); + break; + } + else if (shared_check[adr] == -1) { + old = atomicCAS(shared_check + adr, -1, key); + if (old == -1) { + atomic_fadd(shared_value + adr, aval * bval); + break; + } + } + else { + hash = (hash + 1) & (SH_ROW - 1); + adr = soffset + hash; + } + } + } + } + } + + for (j = tid; j < SH_ROW; j += WARP) { + if (shared_check[soffset + j] != -1) { + index = atomicAdd(d_nz + rid, 1); + shared_check[soffset + index] = shared_check[soffset + j]; + shared_value[soffset + index] = shared_value[soffset + j]; + } + } + int nz = d_nz[rid]; + /* Sorting for shared data */ + int count, target; + for (j = tid; j < nz; j += WARP) { + target = shared_check[soffset + j]; + count = 0; + for (k = 0; k < nz; k++) { + count += (unsigned int)(shared_check[soffset + k] - target) >> 31; + } + d_ccol[offset + count] = shared_check[soffset + j]; + d_cval[offset + count] = shared_value[soffset + j]; + } +} + +template +__global__ void calculate_value_col_bin_each_tb(const int *d_arpt, + const int *d_acol, + const real *d_aval, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const real* __restrict__ d_bval, + int *d_crpt, + int *d_ccol, + real *d_cval, + const int *d_row_perm, + int *d_nz, + int bin_offset, + int bin_size) +{ + int rid = blockIdx.x; + int tid = threadIdx.x & (WARP - 1); + int wid = threadIdx.x / WARP; + int wnum = blockDim.x / WARP; + int j; + __shared__ int shared_check[SH_ROW]; + __shared__ real shared_value[SH_ROW]; + + for (j = threadIdx.x; j < SH_ROW; j += blockDim.x) { + shared_check[j] = -1; + shared_value[j] = 0; + } + + if (rid >= bin_size) { + return; + } + + rid = d_row_perm[rid + bin_offset]; + + if (threadIdx.x == 0) { + d_nz[rid] = 0; + } + __syncthreads(); + + int acol; + int k; + int bcol, hash, key; + int offset = d_crpt[rid]; + int old, index; + real aval, bval; + + for (j = d_arpt[rid] + wid; j < d_arpt[rid + 1]; j += wnum) { + acol = ld_gbl_int32(d_acol + j); + aval = ld_gbl_real(d_aval + j); + for (k = d_brpt[acol] + tid; k < d_brpt[acol + 1]; k += WARP) { + bcol = d_bcol[k]; + bval = d_bval[k]; + + key = bcol; + hash = (bcol * HASH_SCAL) & (SH_ROW - 1); + while (1) { + if (shared_check[hash] == key) { + atomic_fadd(shared_value + hash, aval * bval); + break; + } + else if (shared_check[hash] == -1) { + old = atomicCAS(shared_check + hash, -1, key); + if (old == -1) { + atomic_fadd(shared_value + hash, aval * bval); + break; + } + } + else { + hash = (hash + 1) & (SH_ROW - 1); + } + } + } + } + + __syncthreads(); + if (threadIdx.x < WARP) { + for (j = tid; j < SH_ROW; j += WARP) { + if (shared_check[j] != -1) { + index = atomicAdd(d_nz + rid, 1); + shared_check[index] = shared_check[j]; + shared_value[index] = shared_value[j]; + } + } + } + __syncthreads(); + int nz = d_nz[rid]; + /* Sorting for shared data */ + int count, target; + for (j = threadIdx.x; j < nz; j += blockDim.x) { + target = shared_check[j]; + count = 0; + for (k = 0; k < nz; k++) { + count += (unsigned int)(shared_check[k] - target) >> 31; + } + d_ccol[offset + count] = shared_check[j]; + d_cval[offset + count] = shared_value[j]; + } + +} + +__global__ void calculate_value_col_bin_each_gl(const int *d_arpt, + const int *d_acol, + const real *d_aval, + const int* __restrict__ d_brpt, + const int* __restrict__ d_bcol, + const real* __restrict__ d_bval, + int *d_crpt, + int *d_ccol, + real *d_cval, + const int *d_row_perm, + int *d_nz, + int *d_check, + real *d_value, + int max_row_nz, + int bin_offset, + int M) +{ + int rid = blockIdx.x; + int tid = threadIdx.x & (WARP - 1); + int wid = threadIdx.x / WARP; + int wnum = blockDim.x / WARP; + int j; + + if (rid >= M) { + return; + } + + int doffset = rid * max_row_nz; + + rid = d_row_perm[rid + bin_offset]; + + if (threadIdx.x == 0) { + d_nz[rid] = 0; + } + __syncthreads(); + + int acol; + int k; + int bcol, hash, key, adr; + int offset = d_crpt[rid]; + int old, index; + real aval, bval; + + for (j = d_arpt[rid] + wid; j < d_arpt[rid + 1]; j += wnum) { + acol = ld_gbl_int32(d_acol + j); + aval = ld_gbl_real(d_aval + j); + for (k = d_brpt[acol] + tid; k < d_brpt[acol + 1]; k += WARP) { + bcol = d_bcol[k]; + bval = d_bval[k]; + + key = bcol; + hash = (bcol * HASH_SCAL) % max_row_nz; + adr = doffset + hash; + while (1) { + if (d_check[adr] == key) { + atomic_fadd(d_value + adr, aval * bval); + break; + } + else if (d_check[adr] == -1) { + old = atomicCAS(d_check + adr, -1, key); + if (old == -1) { + atomic_fadd(d_value + adr, aval * bval); + break; + } + } + else { + hash = (hash + 1) % max_row_nz; + adr = doffset + hash; + } + } + } + } + + __syncthreads(); + if (threadIdx.x < WARP) { + for (j = tid; j < max_row_nz; j += WARP) { + if (d_check[doffset + j] != -1) { + index = atomicAdd(d_nz + rid, 1); + d_check[doffset + index] = d_check[doffset + j]; + d_value[doffset + index] = d_value[doffset + j]; + } + } + } + __syncthreads(); + int nz = d_nz[rid]; + + /* Sorting for shared data */ + int count, target; + for (j = threadIdx.x; j < nz; j += blockDim.x) { + target = d_check[doffset + j]; + count = 0; + for (k = 0; k < nz; k++) { + count += (unsigned int)(d_check[doffset + k] - target) >> 31; + } + d_ccol[offset + count] = d_check[doffset + j]; + d_cval[offset + count] = d_value[doffset + j]; + } + +} + +void calculate_value_col_bin(int *d_arpt, int *d_acol, real *d_aval, + int *d_brpt, int *d_bcol, real *d_bval, + int *d_crpt, int *d_ccol, real *d_cval, + sfBIN *bin, + int M); + +void set_row_nnz(int *d_arpt, int *d_acol, + int *d_brpt, int *d_bcol, + int *d_crpt, + sfBIN *bin, + int M, + int *nnz) +{ + int i; + int GS, BS; + for (i = BIN_NUM - 1; i >= 0; i--) { + if (bin->bin_size[i] > 0) { + switch (i) { + case 0: + BS = 512; + GS = div_round_up(bin->bin_size[i] * PWARP, BS); + set_row_nz_bin_pwarp<<stream[i]>>> + (d_arpt, d_acol, + d_brpt, d_bcol, + bin->d_row_perm, + bin->d_row_nz, + bin->bin_offset[i], + bin->bin_size[i]); + break; + case 1 : + BS = 64; + GS = bin->bin_size[i]; + set_row_nz_bin_each_tb<512><<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 2 : + BS = 128; + GS = bin->bin_size[i]; + set_row_nz_bin_each_tb<1024><<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 3 : + BS = 256; + GS = bin->bin_size[i]; + set_row_nz_bin_each_tb<2048><<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 4 : + BS = 512; + GS = bin->bin_size[i]; + set_row_nz_bin_each_tb<4096><<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 5 : + BS = 1024; + GS = bin->bin_size[i]; + set_row_nz_bin_each_tb<8192><<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 6 : + { + int fail_count; + fail_count = 0; + int *d_fail_count, *d_fail_perm; + checkCudaErrors(cudaMalloc((void **)&d_fail_count, sizeof(int))); + checkCudaErrors(cudaMalloc((void **)&d_fail_perm, sizeof(int) * bin->bin_size[i])); + cudaMemcpy(d_fail_count, &fail_count, sizeof(int), cudaMemcpyHostToDevice); + BS = 1024; + GS = bin->bin_size[i]; + set_row_nz_bin_each_tb_large<8192><<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + bin->d_row_perm, bin->d_row_nz, + d_fail_count, d_fail_perm, + bin->bin_offset[i], bin->bin_size[i]); + cudaMemcpy(&fail_count, d_fail_count, sizeof(int), cudaMemcpyDeviceToHost); + if (fail_count > 0) { + int max_row_nz = bin->max_intprod; + size_t table_size = (size_t)max_row_nz * fail_count; + int *d_check; + checkCudaErrors(cudaMalloc((void **)&(d_check), sizeof(int) * table_size)); + + BS = 1024; + GS = div_round_up(table_size, BS); + init_check<<stream[i]>>>(d_check, table_size); + GS = bin->bin_size[i]; + set_row_nz_bin_each_gl<<stream[i]>>> + (d_arpt, d_acol, d_brpt, d_bcol, + d_fail_perm, bin->d_row_nz, d_check, + max_row_nz, 0, fail_count); + cudaFree(d_check); + } + cudaFree(d_fail_count); + cudaFree(d_fail_perm); + } + break; + default : + exit(0); + } + } + } + cudaThreadSynchronize(); + + /* Set row pointer of matrix C */ + thrust::exclusive_scan(thrust::device, bin->d_row_nz, bin->d_row_nz + (M + 1), d_crpt, 0); + cudaMemcpy(nnz, d_crpt + M, sizeof(int), cudaMemcpyDeviceToHost); +} + +void calculate_value_col_bin(int *d_arpt, int *d_acol, real *d_aval, + int *d_brpt, int *d_bcol, real *d_bval, + int *d_crpt, int *d_ccol, real *d_cval, + sfBIN *bin, + int M) +{ + int i; + int GS, BS; + for (i = BIN_NUM - 1; i >= 0; i--) { + if (bin->bin_size[i] > 0) { + switch (i) { + case 0: + BS = 512; + GS = div_round_up(bin->bin_size[i] * PWARP, BS); + calculate_value_col_bin_pwarp<<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 1: + BS = 64; + GS = bin->bin_size[i]; + calculate_value_col_bin_each_tb<256><<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 2: + BS = 128; + GS = bin->bin_size[i]; + calculate_value_col_bin_each_tb<512><<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 3: + BS = 256; + GS = bin->bin_size[i]; + calculate_value_col_bin_each_tb<1024><<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 4: + BS = 512; + GS = bin->bin_size[i]; + calculate_value_col_bin_each_tb<2048><<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 5: + BS = 1024; + GS = bin->bin_size[i]; + calculate_value_col_bin_each_tb<4096><<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + break; + case 6 : + { + int max_row_nz = bin->max_nz * 2; + int table_size = max_row_nz * bin->bin_size[i]; + int *d_check; + real *d_value; + checkCudaErrors(cudaMalloc((void **)&(d_check), sizeof(int) * table_size)); + checkCudaErrors(cudaMalloc((void **)&(d_value), sizeof(real) * table_size)); + + BS = 1024; + GS = div_round_up(table_size, BS); + init_check<<stream[i]>>>(d_check, table_size); + init_value<<stream[i]>>>(d_value, table_size); + GS = bin->bin_size[i]; + calculate_value_col_bin_each_gl<<stream[i]>>> + (d_arpt, d_acol, d_aval, + d_brpt, d_bcol, d_bval, + d_crpt, d_ccol, d_cval, + bin->d_row_perm, bin->d_row_nz, + d_check, d_value, max_row_nz, + bin->bin_offset[i], bin->bin_size[i]); + cudaFree(d_check); + cudaFree(d_value); + } + break; + default : + exit(0); + } + } + } + cudaThreadSynchronize(); +} + +__device__ int binary_searchd(double *arr, double val, int imin, int imax) { + + int ans = 0; + while (imax >= imin) { + int imid = (imin + imax) / 2; + // printf("imin: %d imax: %d imid: %d\n", imin, imax, imid); + + if (arr[imid] < val) { + imin = imid + 1; + } else { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +__device__ int binary_searchf(float *arr, float val, int imin, int imax) { + + int ans = 0; + while (imax >= imin) { + int imid = (imin + imax) / 2; + + if (arr[imid] < val) { + imin = imid + 1; + } else { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +__device__ int binary_searchfpf(double *arr, double val, int imin, int imax) { + + int ans = 0; + while (imax >= imin) { + int imid = (imin + imax) / 2; + printf("imid: %d val: %f arr[imid]: %f\n", imid, val, arr[imid]); + + if (arr[imid] < val) { + imin = imid + 1; + } else { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +__device__ long binary_searchl(long *arr, long val, long imin, long imax) { + + long ans = -1; + while (imax >= imin) { + long imid = (imin + imax) / 2; + + if (arr[imid] <= val) { + imin = imid + 1; + } else if (arr[imid] > val) { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +__device__ long binary_searchi(int *arr, int val, int imin, int imax) { + + int ans = -1; + while (imax >= imin) { + int imid = (imin + imax) / 2; + + if (arr[imid] <= val) { + imin = imid + 1; + } else if (arr[imid] > val) { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +// Binary search that returns exact location, not insertion point (for COO row selection) +__device__ long binary_search_rowselect(long *arr, long val, long imin, long imax) { + while (imax >= imin) { + long imid = (imin + imax) / 2; + + if (arr[imid] < val) { + imin = imid + 1; + } else if (arr[imid] > val) { + imax = imid - 1; + } else { + return imid; + } + } + + return imin + 1; +} at::Tensor expand_values_if_needed(const at::Tensor& values) { // expand @@ -53,6 +1440,15 @@ at::Tensor sparse_coo_tensor_gpu(const at::Tensor& indices, sparse_dim, dense_dim, size, indices, values, values.options().layout(at::kSparse)); } +at::Tensor sparse_csr_tensor_gpu(const at::Tensor& crows, + const at::Tensor& cols, + const at::Tensor& values, + at::ArrayRef size) { + + // return at::sparse_compressed_tensor(crows, cols, values, size, c10::Layout::SparseCsr); // torch 1.13 + return at::sparse_csr_tensor(crows, cols, values, size, values.options().layout(at::kSparseCsr)); // 1.11 +} + template void printCusparseDnMat(int64_t rows, int64_t cols, int64_t ld, T *values_dev) { T* values_host = new T[rows*cols]; @@ -118,8 +1514,6 @@ void spmm_gpu(const at::Tensor& A_rowindices, at::Tensor& B, at::Tensor& C) { - // cusparseHandle_t handle; - // CHECK_CUSPARSE(cusparseCreate(&handle)); auto handle = at::cuda::getCurrentCUDASparseHandle(); // Impl1 -- coo2csr + csrmm2 @@ -128,12 +1522,8 @@ void spmm_gpu(const at::Tensor& A_rowindices, clock_t start, stop; int32_t *d_a_csrrows; - - // int devid_old = 0; - // cudaGetDevice(&devid_old); - // cudaSetDevice(devid); - cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); + CHECK_CUSPARSE(cusparseXcoo2csr(handle, A_rowindices.data(), nnz, @@ -151,8 +1541,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, cusparseSpMatDescr_t matA; CHECK_CUSPARSE(cusparseCreateCsr(&matA, n, // rows - // b_col, // cols - m, // cols + m, // cols nnz, // nnz d_a_csrrows, // csrRowOffsets A_colindices.data(), // csrColInd @@ -162,23 +1551,14 @@ void spmm_gpu(const at::Tensor& A_rowindices, CUSPARSE_INDEX_BASE_ZERO, // idxBase, CUDA_R_32F)); // valueType - // Row-major to column-major - B.t_(); - B.set_data(B.contiguous()); - B.set_data(B.view({b_row, b_col})); - cusparseDnMatDescr_t matB; CHECK_CUSPARSE(cusparseCreateDnMat(&matB, - // b_col, // rows - b_row, // rows - // b_row, // cols - b_col, // cols - // b_col, // ld - b_row, // ld + b_col, // rows + b_row, // cols + b_col, // ld B.data(), // values CUDA_R_32F, // valueType CUSPARSE_ORDER_COL)); // order - // Row-major to column-major C.t_(); C.set_data(C.contiguous()); @@ -187,8 +1567,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, cusparseDnMatDescr_t matC; CHECK_CUSPARSE(cusparseCreateDnMat(&matC, n, // rows - B.size(1), // cols - // n, // ld + b_col, // cols n, // ld C.data(), // values CUDA_R_32F, // valueType @@ -197,49 +1576,1447 @@ void spmm_gpu(const at::Tensor& A_rowindices, size_t bufferSize; CHECK_CUSPARSE(cusparseSpMM_bufferSize(handle, // handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA - // CUSPARSE_OPERATION_TRANSPOSE, // opB - CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + CUSPARSE_OPERATION_TRANSPOSE, // opB &alpha, // alpha matA, // matA matB, // matB &beta, // beta matC, // matC CUDA_R_32F, // computeType - CUSPARSE_CSRMM_ALG1, // alg + // CUSPARSE_CSRMM_ALG1, // alg + CUSPARSE_SPMM_CSR_ALG1, // alg &bufferSize)); // bufferSize - void* d_buffer = NULL; cudaMalloc(&d_buffer, bufferSize); CHECK_CUSPARSE(cusparseSpMM(handle, // handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA - // CUSPARSE_OPERATION_TRANSPOSE, // opB - CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + CUSPARSE_OPERATION_TRANSPOSE, // opB &alpha, // alpha matA, // matA matB, // matB &beta, // beta matC, // matC CUDA_R_32F, // computeType - CUSPARSE_CSRMM_ALG1, // alg + // CUSPARSE_CSRMM_ALG1, // alg + CUSPARSE_SPMM_CSR_ALG1, // alg d_buffer)); // buffer cudaFree(d_a_csrrows); cudaFree(d_buffer); - CHECK_ERROR("spmm_gpu error") // Column-major to row-major C.set_data(C.view({c_col, c_row})); C.t_(); +} +std::vector coogeam_gpu( + const at::Tensor& A_rowindices, + const at::Tensor& A_colindices, + const at::Tensor& A_values, + const at::Tensor& B_rowindices, + const at::Tensor& B_colindices, + const at::Tensor& B_values, + int32_t m, + int32_t n) { - // Column-major to row-major - B.set_data(B.view({b_col, b_row})); - B.t_(); + + // sum matrix A (m x n) and B (m x n) + auto handle = at::cuda::getCurrentCUDASparseHandle(); + + int32_t *d_a_csrrows = NULL; + int32_t *d_b_csrrows = NULL; + + int32_t A_nnz = A_values.size(0); + int32_t B_nnz = B_values.size(0); + + + // Construct CSR offsets array for A and B + cudaMalloc(&d_a_csrrows, (m + 1) * sizeof(int32_t)); + CHECK_CUSPARSE(cusparseXcoo2csr(handle, + A_rowindices.data(), + A_nnz, + m, + d_a_csrrows, + CUSPARSE_INDEX_BASE_ZERO)); + + cudaMalloc(&d_b_csrrows, (m + 1) * sizeof(int32_t)); + CHECK_CUSPARSE(cusparseXcoo2csr(handle, + B_rowindices.data(), + B_nnz, + m, + d_b_csrrows, + CUSPARSE_INDEX_BASE_ZERO)); + + double alpha = 1; + double beta = 1; + + // Construct CSR matrix structs for A, B, and C (empty CSR) + cusparseMatDescr_t matA; + cusparseCreateMatDescr(&matA); + + cusparseMatDescr_t matB; + cusparseCreateMatDescr(&matB); + + int32_t *d_c_csrrows = NULL; + int32_t *d_c_columns = NULL; + double *d_c_values = NULL; + + cudaMalloc(&d_c_csrrows, (m + 1) * sizeof(int32_t)); + cusparseMatDescr_t matC; + cusparseCreateMatDescr(&matC); + + char *buffer = NULL; + size_t bufferSize; + CHECK_CUSPARSE(cusparseDcsrgeam2_bufferSizeExt(handle, // handle, + m, // m, + n, // n, + &alpha, // alpha, + matA, // descrA, + A_nnz, // nnzA, + A_values.data(), // csrSortedValA, + d_a_csrrows, // csrSortedRowPtrA, + A_colindices.data(), // csrSortedColIndA, + &beta, // beta, + matB, // descrB, + B_nnz, // nnzB, + B_values.data(), // csrSortedValB, + d_b_csrrows, // csrSortedRowPtrB, + B_colindices.data(), // csrSortedColIndB, + matC, // descrC, + d_c_values, // csrSortedValC, + d_c_csrrows, // csrSortedRowPtrC, + d_c_columns, // csrSortedColIndC, + &bufferSize)); // pBufferSizeInBytes) + + cudaMalloc(&buffer, bufferSize * sizeof(char)); + + int32_t C_nnz; + int32_t *C_nnz_ptr = &C_nnz; + // std::cout << "m: " << m << std::endl; + // std::cout << "n: " << n << std::endl; + // std::cout << "A_nnz: " << A_nnz << std::endl; + // std::cout << "B_nnz: " << B_nnz << std::endl; + // std::cout << "d_a_csrrows: " << d_a_csrrows << std::endl; + // std::cout << "d_b_csrrows: " << d_b_csrrows << std::endl; + // std::cout << "d_c_csrrows: " << d_c_csrrows << std::endl; + // std::cout << "A_colindices.size(0): " << A_colindices.size(0) << std::endl; + // std::cout << "B_colindices.size(0): " << B_colindices.size(0) << std::endl; + CHECK_CUSPARSE(cusparseXcsrgeam2Nnz(handle, // handle, + m, // m, + n, // n, + matA, // descrA, + A_nnz, // nnzA, + d_a_csrrows, // csrSortedRowPtrA, + A_colindices.data(), // csrSortedColIndA, + matB, // descrB, + B_nnz, // nnzB, + d_b_csrrows, // csrSortedRowPtrB, + B_colindices.data(), // csrSortedColIndB, + matC, // descrC, + d_c_csrrows, // csrSortedRowPtrC, + C_nnz_ptr, // nnzTotalDevHostPtr, + buffer)); // workspace + if (NULL != C_nnz_ptr) { + C_nnz = *C_nnz_ptr; + } else { + int32_t baseC; + cudaMemcpy(&C_nnz, d_c_csrrows + m, sizeof(int32_t), cudaMemcpyDeviceToHost); + cudaMemcpy(&baseC, d_c_csrrows, sizeof(int32_t), cudaMemcpyDeviceToHost); + C_nnz -= baseC; + } + + cudaMalloc(&d_c_columns, C_nnz * sizeof(int32_t)); + cudaMalloc(&d_c_values, C_nnz * sizeof(double)); + + CHECK_CUSPARSE(cusparseDcsrgeam2(handle, // handle, + m, // m, + n, // n, + &alpha, // alpha, + matA, // descrA, + A_nnz, // nnzA, + A_values.data(), // csrSortedValA, + d_a_csrrows, // csrSortedRowPtrA, + A_colindices.data(), // csrSortedColIndA, + &beta, // beta, + matB, // descrB, + B_nnz, // nnzB, + B_values.data(), // csrSortedValB, + d_b_csrrows, // csrSortedRowPtrB, + B_colindices.data(), // csrSortedColIndB, + matC, // descrC, + d_c_values, // csrSortedValC, + d_c_csrrows, // csrSortedRowPtrC, + d_c_columns, // csrSortedColIndC, + buffer)); // pBuffer) + + auto rows_dim = torch::IntArrayRef{m + 1}; + auto nnz_dim = torch::IntArrayRef{C_nnz}; + + auto c_csrOffsets = torch::from_blob(d_c_csrrows, rows_dim, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + auto c_columns = torch::from_blob(d_c_columns, nnz_dim, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + auto c_values = torch::from_blob(d_c_values, nnz_dim, torch::TensorOptions() + .dtype(torch::kFloat64) + .device(torch::kCUDA)); + cudaFree(buffer); + + CHECK_CUSPARSE(cusparseDestroyMatDescr(matA)); + CHECK_CUSPARSE(cusparseDestroyMatDescr(matB)); + CHECK_CUSPARSE(cusparseDestroyMatDescr(matC)); + + return {c_csrOffsets, c_columns, c_values}; +} + +std::vector nsparse_spgemm( + const at::Tensor& A_crow_indices, + const at::Tensor& A_col_indices, + const at::Tensor& A_values, + const at::Tensor& B_crow_indices, + const at::Tensor& B_col_indices, + const at::Tensor& B_values, + int32_t n, + int32_t k, + int32_t m) { + + sfBIN bin; + + /* Set max bin */ + int *d_a_crow_indices = A_crow_indices.data(); + int *d_a_col_indices = A_col_indices.data(); + float *d_a_values = A_values.data(); + + int *d_b_crow_indices = B_crow_indices.data(); + int *d_b_col_indices = B_col_indices.data(); + float *d_b_values = B_values.data(); + + init_bin(&bin, n); + std::cout << "after init_bin" << std::endl; + + // int *d_c_crow_indices = C_crow_indices.data(); + // int *d_c_crow_indices; + // cudaMalloc(&d_c_crow_indices, sizeof(int) * (n + 1)); + torch::Tensor C_crow_indices = torch::zeros({n + 1}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + int *d_c_crow_indices = C_crow_indices.data(); + + set_max_bin(d_a_crow_indices, d_a_col_indices, d_b_crow_indices, &bin, n); + std::cout << "after set_max_bin" << std::endl; + + int C_nnz; + /* Count nz of C */ + set_row_nnz(d_a_crow_indices, d_a_col_indices, + d_b_crow_indices, d_b_col_indices, + d_c_crow_indices, + &bin, + n, + &C_nnz); + std::cout << "after set_row_nnz" << std::endl; + + /* Set bin */ + set_min_bin(&bin, n); + std::cout << "after set_min_bin" << std::endl; + + // int *d_c_col_indices; + // float *d_c_values; + // cudaMalloc(&d_c_col_indices, sizeof(int) * C_nnz); + // cudaMalloc(&d_c_values, sizeof(float) * C_nnz); + torch::Tensor C_col_indices = torch::zeros({C_nnz}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + + torch::Tensor C_values = torch::zeros({C_nnz}, torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA)); + int *d_c_col_indices = C_col_indices.data(); + float *d_c_values = C_values.data(); + + /* Calculating value of C */ + calculate_value_col_bin(d_a_crow_indices, d_a_col_indices, d_a_values, + d_b_crow_indices, d_b_col_indices, d_b_values, + d_c_crow_indices, d_c_col_indices, d_c_values, + &bin, + n); + std::cout << "after calculate_value_col_bin" << std::endl; + + release_bin(bin); + std::cout << "after release_bin" << std::endl; + + CHECK_ERROR("nsparse_spgemm") + + return {C_crow_indices, C_col_indices, C_values}; + // return {c_csr_offsets, c_columns, c_values}; +} + +std::vector spgemm_gpu( + const at::Tensor& A_crow_indices, + const at::Tensor& A_col_indices, + const at::Tensor& A_values, + const at::Tensor& B_crow_indices, + const at::Tensor& B_col_indices, + const at::Tensor& B_values, + const at::Tensor& C_crow_indices, + const at::Tensor& C_col_indices, + const at::Tensor& C_values, + int32_t n, + int32_t k, + int32_t m) { + + + // SpGEMM between A (n x k) and B (k x m) + auto handle = at::cuda::getCurrentCUDASparseHandle(); + + int32_t A_nnz = A_values.size(0); + int32_t B_nnz = B_values.size(0); + int32_t C_nnz = C_values.size(0); + + float alpha = 1; + float beta = 1; + + // Construct CSR matrix structs for A, B, and C (empty CSR) + cusparseSpMatDescr_t matA; + CHECK_CUSPARSE(cusparseCreateCsr(&matA, + n, // rows + k, // cols + A_nnz, // nnz + A_crow_indices.data(), // csrRowOffsets + A_col_indices.data(), // csrColInd + A_values.data(), // csrValues + CUSPARSE_INDEX_32I, // csrRowOffsetsType + CUSPARSE_INDEX_32I, // csrColIndType + CUSPARSE_INDEX_BASE_ZERO, // idxBase, + CUDA_R_32F)); // valueType + + cusparseSpMatDescr_t matB; + CHECK_CUSPARSE(cusparseCreateCsr(&matB, + k, // rows + m, // cols + B_nnz, // nnz + B_crow_indices.data(), // csrRowOffsets + B_col_indices.data(), // csrColInd + B_values.data(), // csrValues + CUSPARSE_INDEX_32I, // csrRowOffsetsType + CUSPARSE_INDEX_32I, // csrColIndType + CUSPARSE_INDEX_BASE_ZERO, // idxBase, + CUDA_R_32F)); // valueType + + cusparseSpMatDescr_t matC; + CHECK_CUSPARSE(cusparseCreateCsr(&matC, + n, // rows + m, // cols + C_nnz, // nnz + C_crow_indices.data(), // csrRowOffsets + C_col_indices.data(), // csrColInd + C_values.data(), // csrValues + CUSPARSE_INDEX_32I, // csrRowOffsetsType + CUSPARSE_INDEX_32I, // csrColIndType + CUSPARSE_INDEX_BASE_ZERO, // idxBase, + CUDA_R_32F)); // valueType + + cusparseSpGEMMDescr_t spgemmDescr; + CHECK_CUSPARSE(cusparseSpGEMM_createDescr(&spgemmDescr)); + + void *dBuffer1 = NULL; + void *dBuffer2 = NULL; + size_t bufferSize1 = 0; + size_t bufferSize2 = 0; + CHECK_CUSPARSE(cusparseSpGEMM_workEstimation(handle, // handle + CUSPARSE_OPERATION_NON_TRANSPOSE, // opA + CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize1, // bufferSize1 + NULL)); // externalBuffer1 + + //CHECK_ERROR(cudaMalloc(&dBuffer1, bufferSize1)); + cudaMalloc(&dBuffer1, bufferSize1); + + CHECK_CUSPARSE(cusparseSpGEMM_workEstimation(handle, // handle + CUSPARSE_OPERATION_NON_TRANSPOSE, // opA + CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize1, // bufferSize1 + dBuffer1)); // externalBuffer1 + + CHECK_CUSPARSE(cusparseSpGEMM_compute(handle, + CUSPARSE_OPERATION_NON_TRANSPOSE, // opA + CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize2, // bufferSize2 + NULL)); // externalBuffer2 + + cudaMalloc(&dBuffer2, bufferSize2); + + CHECK_CUSPARSE(cusparseSpGEMM_compute(handle, + CUSPARSE_OPERATION_NON_TRANSPOSE, // opA + CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize2, // bufferSize2 + dBuffer2)); // externalBuffer2 + + int64_t C_num_rows1, C_num_cols1, C_nnz1; + CHECK_CUSPARSE(cusparseSpMatGetSize(matC, &C_num_rows1, &C_num_cols1, &C_nnz1)); + + int32_t *dC_csrOffsets; + int32_t *dC_columns; + float *dC_values; + cudaMalloc(&dC_csrOffsets, (n + 1) * sizeof(int32_t)); + cudaMalloc(&dC_columns, C_nnz1 * sizeof(int32_t)); + cudaMalloc(&dC_values, C_nnz1 * sizeof(float)); + + CHECK_CUSPARSE(cusparseCsrSetPointers(matC, dC_csrOffsets, dC_columns, dC_values)); + CHECK_CUSPARSE(cusparseSpGEMM_copy(handle, + CUSPARSE_OPERATION_NON_TRANSPOSE, // opA + CUSPARSE_OPERATION_NON_TRANSPOSE, // opB + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr)); + + auto rows_dim = torch::IntArrayRef{n + 1}; + auto nnz_dim = torch::IntArrayRef{C_nnz1}; + + auto c_csrOffsets = torch::from_blob(dC_csrOffsets, rows_dim, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + auto c_columns = torch::from_blob(dC_columns, nnz_dim, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + auto c_values = torch::from_blob(dC_values, nnz_dim, torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA)); + + CHECK_ERROR("spgemm_gpu error"); + cudaFree(dBuffer1); + cudaFree(dBuffer2); + + CHECK_CUSPARSE(cusparseSpGEMM_destroyDescr(spgemmDescr)); + CHECK_CUSPARSE(cusparseDestroySpMat(matA)); + CHECK_CUSPARSE(cusparseDestroySpMat(matB)); + CHECK_CUSPARSE(cusparseDestroySpMat(matC)); + + return {c_csrOffsets, c_columns, c_values}; +} + +__global__ void DownSample(long *h_counts, long *h_rows, long *ps_h_rows, long *hev_indices, int *overflow, + int nnz) { + + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (long i = id; i < nnz; i += stride) { + if (i - ps_h_rows[h_rows[i]] < overflow[h_rows[i]]) { + h_counts[hev_indices[i]] = 0; + } + } +} + +void downsample_gpu(const at::Tensor& h_counts, + const at::Tensor& h_rows, + const at::Tensor& ps_h_rows, + const at::Tensor& hev_indices, + const at::Tensor& overflow, + int nnz) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(nnz / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + DownSample<<>>(h_counts.data(), + h_rows.data(), + ps_h_rows.data(), + hev_indices.data(), + overflow.data(), + nnz); + CHECK_ERROR("downsampling error") +} + +__global__ void ComputeDarts(float *dartx_values, float *darty_values, long *neighbor_sizes, + long *psum_neighbor_sizes, float *pmax, int n_darts, int mb_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + int dart_count = n_darts * mb_count; + for (int i = id; i < dart_count; i += stride) { + long row = i / n_darts; + dartx_values[i] *= neighbor_sizes[row]; + dartx_values[i] += psum_neighbor_sizes[row]; + darty_values[i] *= pmax[row]; + } +} + +void compute_darts_gpu(const at::Tensor& dartx_values, + const at::Tensor& darty_values, + const at::Tensor& neighbor_sizes, + const at::Tensor& psum_neighbor_sizes, + const at::Tensor& pmax, + int n_darts, + int mb_count) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil((n_darts * mb_count) / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ComputeDarts<<>>(dartx_values.data(), + darty_values.data(), + neighbor_sizes.data(), + psum_neighbor_sizes.data(), + pmax.data(), + n_darts, + mb_count); + CHECK_ERROR("dart computation error") +} + +__global__ void ComputeDarts1D(double *dart_values, double *p_rowsum, double *ps_p_rowsum, +// __global__ void ComputeDarts1D(float *dart_values, float *p_rowsum, float *ps_p_rowsum, + long *ps_dart_count_rows, int n_darts, int mb_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + // int dart_count = n_darts * mb_count; + for (int i = id; i < n_darts; i += stride) { + // int row = i / n_darts; + long row = binary_searchl(ps_dart_count_rows, i, 0, mb_count); + dart_values[i] *= p_rowsum[row]; + dart_values[i] += ps_p_rowsum[row]; + } +} + +void compute_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& p_rowsum, + const at::Tensor& ps_p_rowsum, const at::Tensor& ps_dart_count_rows, + int n_darts, int mb_count) { + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(n_darts / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ComputeDarts1D<<>>(dart_values.data(), + p_rowsum.data(), + ps_p_rowsum.data(), + // ComputeDarts1D<<>>(dart_values.data(), + // p_rowsum.data(), + // ps_p_rowsum.data(), + ps_dart_count_rows.data(), + n_darts, + mb_count); + CHECK_ERROR("dart1d computation error") +} + +__global__ void ComputeDartsSelect(double *dart_select, double *dart_hits_inv_sum, double *ps_dart_hits_inv_sum, + long *ps_overflow, long mb_count, long total_overflow) { + + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (long i = id; i < total_overflow; i += stride) { + long row = binary_searchl(ps_overflow, i, 0, mb_count); + dart_select[i] *= dart_hits_inv_sum[row]; + dart_select[i] += ps_dart_hits_inv_sum[row]; + } +} + +void compute_darts_select_gpu(const at::Tensor& dart_select, + const at::Tensor& dart_hits_inv_sum, + const at::Tensor& ps_dart_hits_inv_sum, + const at::Tensor& ps_overflow, + long mb_count, + long total_overflow) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (total_overflow == 0) { + return; + } + + ComputeDartsSelect<<>>(dart_select.data(), + dart_hits_inv_sum.data(), + ps_dart_hits_inv_sum.data(), + ps_overflow.data(), + mb_count, + total_overflow); + CHECK_ERROR("selection dart computation error") +} + +__global__ void ThrowDarts(float *dartx_values, float *darty_values, float *p_values, + long *h_values, int n_darts, int mb_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + int dart_count = n_darts * mb_count; + for (int i = id; i < dart_count; i += stride) { + long dartx_val = (long) dartx_values[i]; + if (darty_values[i] < p_values[dartx_val]) { + atomicAdd((unsigned long long *)&h_values[dartx_val], 1L); + } + } +} + +void throw_darts_gpu(const at::Tensor& dartx_values, + const at::Tensor& darty_values, + const at::Tensor& p_values, + const at::Tensor& h_values, + int n_darts, + int mb_count) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil((n_darts * mb_count) / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ThrowDarts<<>>(dartx_values.data(), + darty_values.data(), + p_values.data(), + h_values.data(), + n_darts, + mb_count); + CHECK_ERROR("dart throwing error") +} + +__global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_values, + int dart_count, int nnz) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < dart_count; i += stride) { + // int vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz - 1); + int vtx = binary_searchd(ps_p_values, dart_values[i], 0, nnz - 1); + // printf("i: %d vtx: %d\n", i, vtx); + // if (vtx < 0 || vtx >= nnz) { + // printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); + // } + atomicAdd(&h_values[vtx], 1); + } +} + +void throw_darts1d_gpu(const at::Tensor& dart_values, + const at::Tensor& ps_p_values, + const at::Tensor& h_values, + int dart_count, + int nnz) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil((dart_count) / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + // ThrowDarts1D<<>>(dart_values.data(), + // ps_p_values.data(), + ThrowDarts1D<<>>(dart_values.data(), + ps_p_values.data(), + h_values.data(), + dart_count, + nnz); + CHECK_ERROR("dart throwing error") +} + +__global__ void ThrowDartsSelect(double *dart_select, double *ps_dart_hits_inv, int *dart_hits_count, + int total_overflow, int nnz) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < total_overflow; i += stride) { + // int vtx = binary_searchf(ps_dart_hits_inv, dart_select[i], 0, nnz); + int vtx = binary_searchd(ps_dart_hits_inv, dart_select[i], 0, nnz); + atomicAnd(&dart_hits_count[vtx], 0); + } +} + +void throw_darts_select_gpu(const at::Tensor& dart_select, + const at::Tensor& ps_dart_hits_inv, + const at::Tensor& dart_hits_count, + int total_overflow, + int nnz) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (total_overflow == 0) { + return; + } + + ThrowDartsSelect<<>>(dart_select.data(), + ps_dart_hits_inv.data(), + dart_hits_count.data(), + total_overflow, + nnz); + CHECK_ERROR("selection dart throwing error") +} + +__global__ void Normalized(double *output, double *input, long *index, int len) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < len; i += stride) { + output[i] /= input[index[i]]; + } +} + +void normalized_gpu(const at::Tensor& output, const at::Tensor& input, const at::Tensor& index, int len) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + Normalized<<>>(output.data(), input.data(), index.data(), len); +} + +__global__ void Normalize(double *output, double *input, long *index, int len) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < len; i += stride) { + output[i] /= input[index[i]]; + } +} + +void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at::Tensor& index, int len) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + Normalize<<>>(output.data(), input.data(), index.data(), len); +} + +__global__ void NormalizeCsr(double *values, long *crows, int len) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < len; i += stride) { + // output[i] /= input[index[i]]; + int sum = 0; + for (int j = crows[i]; j < crows[i + 1]; j++) { + sum += values[j]; + } + for (int j = crows[i]; j < crows[i + 1]; j++) { + values[j] /= sum; + } + } +} + +void normalize_csr_gpu(const at::Tensor& values, const at::Tensor& crows, int len) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + NormalizeCsr<<>>(values.data(), crows.data(), len); +} + +__global__ void SumCsrd(double *output, long *crows, double *values, int len) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < len; i += stride) { + // output[i] /= input[index[i]]; + for (int j = crows[i]; j < crows[i + 1]; j++) { + output[i] += values[j]; + } + } +} + +void sum_csrd_gpu(const at::Tensor& output, const at::Tensor& crows, const at::Tensor& values, int len) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + SumCsrd<<>>(output.data(), crows.data(), values.data(), len); +} + +__global__ void SumCsri(int *output, long *crows, int *values, int len) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < len; i += stride) { + // output[i] /= input[index[i]]; + for (int j = crows[i]; j < crows[i + 1]; j++) { + output[i] += values[j]; + } + } +} + +void sum_csri_gpu(const at::Tensor& output, const at::Tensor& crows, const at::Tensor& values, int len) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + SumCsri<<>>(output.data(), crows.data(), values.data(), len); +} + +__global__ void NormalizeBatch(float *output, int *input, int output_rows, int output_cols) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < output_rows * output_cols; i += stride) { + int row = i / output_cols; + output[i] /= input[row]; + } +} + +void normalize_batch_gpu(const at::Tensor& output, const at::Tensor& input, int rows, int cols) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil((rows * cols) / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + NormalizeBatch<<>>(output.data(), input.data(), rows, cols); + CHECK_ERROR("normalize batch error") +} + +__global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int rank, int size, + int replication, int nnz, int batch_size, int node_count, int mb_count, + int shift_size, int semibulk_size) { + + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + int rank_c = rank / replication; + long proc_row_chunk = rank_c * ((batch_size * mb_count) / (size / replication)); + + for (int i = id; i < nnz; i += stride) { + long mb_row = (row_select_rows[i] + proc_row_chunk) / shift_size; + // row_shift[i] += mb_row * node_count; + row_shift[i] += (mb_row % semibulk_size) * node_count; + } +} + +void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_select_rows, + int nnz, int rank, int size, int replication, int batch_size, int node_count, + int mb_count, int shift_size, int semibulk_size) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(nnz / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ShiftRowSelect<<>>(row_shift.data(), + row_select_rows.data(), + rank, + size, + replication, + nnz, + batch_size, + node_count, + mb_count, + shift_size, + semibulk_size); + CHECK_ERROR("shift row select error") +} + +__global__ void ShiftColSelect(long *col_shift, int nnz, int batch_size, int node_count) { + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (int i = id; i < nnz; i += stride) { + long mb_row = i / batch_size; + col_shift[i] += mb_row * node_count; + } +} + +void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, int node_count) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(nnz / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ShiftColSelect<<>>(col_shift.data(), nnz, batch_size, node_count); + CHECK_ERROR("shift col select error") +} + +// __global__ void ScatterAddD(double *src, long *indices, double *values, int num_vals) { +// int id = blockIdx.x * blockDim.x + threadIdx.x; +// int stride = blockDim.x * gridDim.x; +// +// for (int i = id; i < num_vals; i += stride) { +// atomicAdd(&src[indices[i]], values[i]); +// } +// } +// +// void scatterd_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at::Tensor& values, int num_vals) { +// +// +// int BLOCK_SIZE = 256; +// int BLOCK_COUNT = std::ceil(num_vals / ((float) BLOCK_SIZE)); +// BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); +// +// ScatterAddD<<>>(src.data(), +// indices.data(), +// values.data(), +// num_vals); +// CHECK_ERROR("scatter add doubles error") +// } + +__global__ void ScatterAddI(int *src, long *indices, int *values, int num_vals) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < num_vals; i += stride) { + atomicAdd(&src[indices[i]], values[i]); + } +} + +void scatteri_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at::Tensor& values, int num_vals) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(num_vals / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ScatterAddI<<>>(src.data(), + indices.data(), + values.data(), + num_vals); + CHECK_ERROR("scatter add ints error") +} + +// // Per-process rowselect coo +// __global__ void RowSelectCoo(long *nnz_cols, long *row_ids, bool *mask, int nnz_col_count, int row_count) { +// int id = blockIdx.x * blockDim.x + threadIdx.x; +// int stride = blockDim.x * gridDim.x; +// +// for (int i = id; i < row_count; i += stride) { +// long idx = binary_search_rowselect(nnz_cols, row_ids[i], 0L, nnz_col_count); +// if (nnz_cols[idx] == row_ids[i]) { +// mask[i] = true; +// } +// } +// } +// +// void rowselect_coo_gpu(const at::Tensor& nnz_cols, const at::Tensor& rows, const at::Tensor& mask, +// int nnz_col_count, int row_count) { +// +// +// int BLOCK_SIZE = 256; +// int BLOCK_COUNT = std::ceil(row_count / ((float) BLOCK_SIZE)); +// BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); +// +// if (nnz_col_count == 0 || row_count == 0) { +// return; +// } +// +// RowSelectCoo<<>>(nnz_cols.data(), +// rows.data(), +// mask.data(), +// nnz_col_count, +// row_count); +// CHECK_ERROR("rowselect coo error") +// } + +// Bulk rowselect coo +__global__ void RowSelectCoo(long **nnz_cols, long *row_ids, bool *mask, int *nnz_cols_counts, int row_count, + int proc_count) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + int mask_len = row_count * proc_count; + for (int i = id; i < mask_len; i += stride) { + long proc = ((long) i) / row_count; + long col = ((long) i) % row_count; + if (nnz_cols_counts[proc] > 0) { + long idx = binary_search_rowselect(nnz_cols[proc], row_ids[col], 0L, (long) nnz_cols_counts[proc]); + if (nnz_cols[proc][idx] == row_ids[col]) { + mask[i] = true; + } + } + } +} + +void rowselect_coo_gpu(std::vector nnz_cols, const at::Tensor& rows, + const at::Tensor& mask, const at::Tensor& nnz_cols_counts, + int row_count, int proc_count) { + + + int mask_len = row_count * proc_count; + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(mask_len / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (mask_len == 0) { + return; + } + + std::vector nnz_cols_hptr(nnz_cols.size()); + for (int i = 0; i < nnz_cols.size(); i++) { + nnz_cols_hptr[i] = nnz_cols[i].data(); + } + + long **nnz_cols_dptr; + cudaMalloc(&nnz_cols_dptr, nnz_cols.size() * sizeof(long*)); + cudaMemcpy(nnz_cols_dptr, nnz_cols_hptr.data(), nnz_cols.size() * sizeof(long*), cudaMemcpyHostToDevice); + + RowSelectCoo<<>>(nnz_cols_dptr, + rows.data(), + mask.data(), + nnz_cols_counts.data(), + row_count, + proc_count); + CHECK_ERROR("rowselect coo error") +} + +// Per-process rowselect csr +__global__ void RowSelectCsr(long *nnz_cols, long *row_offsets, bool *mask, int nnz_col_count) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < nnz_col_count; i += stride) { + long vtx = nnz_cols[i]; + long degree = row_offsets[vtx + 1] - row_offsets[vtx]; + // might be able to use memset here + for (int j = 0; j < degree; j++) { + mask[row_offsets[vtx] + j] = true; + } + } +} + +void rowselect_csr_gpu(const at::Tensor& nnz_cols, const at::Tensor& row_offsets, const at::Tensor& mask, + int nnz_col_count, int nnz_count) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(nnz_col_count / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (nnz_col_count == 0 || nnz_count == 0) { + return; + } + + RowSelectCsr<<>>(nnz_cols.data(), + row_offsets.data(), + mask.data(), + nnz_col_count); + CHECK_ERROR("rowselect csr error") +} + +__global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int nodes_per_proc, int proc_count) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < vtxs_count; i += stride) { + int vtx = (int) vtxs[i]; + int dst_proc = vtx / nodes_per_proc; + if (dst_proc >= proc_count) { + dst_proc = proc_count - 1; + } + atomicAdd(&proc_tally[dst_proc], 1); + } +} + +__global__ void SortVtxs(long *vtxs, long *src_vtx_send, int *ps_proc_tally, long *og_idx, int vtxs_count, + int nodes_per_proc, int proc_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < vtxs_count; i += stride) { + long vtx = vtxs[i]; + int dst_proc = vtx / nodes_per_proc; + if (dst_proc >= proc_count) { + dst_proc = proc_count - 1; + } + int idx = atomicAdd(&ps_proc_tally[dst_proc], 1); + src_vtx_send[idx] = vtx; + og_idx[idx] = i; + } +} + +void sort_dst_proc_gpu(const at::Tensor& vtxs, const at::Tensor& src_vtx_sort, const at::Tensor& og_idxs, + const at::Tensor& tally, int nodes_per_proc, int proc_count) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(vtxs.sizes()[0] / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (nodes_per_proc == 0 || vtxs.sizes()[0] == 0) { + return; + } + + VtxTally<<>>(tally.data(), + vtxs.data(), + vtxs.sizes()[0], + nodes_per_proc, + proc_count); + + auto ps_proc_tally = tally.cumsum(0, torch::kInt32).roll(1); + ps_proc_tally[0] = 0; + + SortVtxs<<>>(vtxs.data(), + src_vtx_sort.data(), + ps_proc_tally.data(), + og_idxs.data(), + vtxs.sizes()[0], + nodes_per_proc, + proc_count); + CHECK_ERROR("sort_dst_proc error") +} + +__global__ void RearrangelRows(long *mata_rows, long *mata_cols, long *matc_crows, long *matb_crows, + long *matb_cols, long *matc_cols, int nnz_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < nnz_count; i += stride) { + long row = mata_rows[i]; + long col = mata_cols[i]; + long dst_idx = matc_crows[row]; + long src_idx = matb_crows[col]; + long row_len = matb_crows[col + 1] - matb_crows[col]; + // memcpy(&matc_cols[dst_idx], &matb_cols[src_idx], row_len * sizeof(long)); + for (int j = 0; j < row_len; j++) { + matc_cols[dst_idx + j] = matb_cols[src_idx + j]; + } + } +} + +void rearrangel_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols, + const at::Tensor& matc_crows, const at::Tensor& matb_crows, + const at::Tensor& matb_cols, const at::Tensor& matc_cols) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(mata_rows.sizes()[0] / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (mata_rows.sizes()[0] == 0) { + return; + } + + RearrangelRows<<>>(mata_rows.data(), + mata_cols.data(), + matc_crows.data(), + matb_crows.data(), + matb_cols.data(), + matc_cols.data(), + mata_rows.sizes()[0]); + + CHECK_ERROR("rearrange_rows error") +} + +__global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, int *matb_crows, + int *matb_cols, int *matc_cols, long nnz_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < nnz_count; i += stride) { + long row = mata_rows[i]; + long col = mata_cols[i]; + long dst_idx = matc_crows[row]; + long src_idx = matb_crows[col]; + long row_len = matb_crows[col + 1] - matb_crows[col]; + // memcpy(&matc_cols[dst_idx], &matb_cols[src_idx], row_len * sizeof(long)); + for (int j = 0; j < row_len; j++) { + matc_cols[dst_idx + j] = (int)(matb_cols[src_idx + j]); + } + } +} + +void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols, + const at::Tensor& matc_crows, const at::Tensor& matb_crows, + const at::Tensor& matb_cols, const at::Tensor& matc_cols) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(mata_rows.sizes()[0] / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (mata_rows.sizes()[0] == 0) { + return; + } + + RearrangeRows<<>>(mata_rows.data(), + mata_cols.data(), + matc_crows.data(), + matb_crows.data(), + matb_cols.data(), + matc_cols.data(), + mata_rows.sizes()[0]); + + CHECK_ERROR("rearrange_rows error") +} + +// __global__ void ReduceSum(long *matc_crows, long *mata_crows, long *matb_crows, long *matc_cols, +// long *mata_cols, long *matb_cols, int row_count) { +// +// int id = blockIdx.x * blockDim.x + threadIdx.x; +// int stride = blockDim.x * gridDim.x; +// +// for (int i = id; i < row_count; i += stride) { +// int row_start = matc_crows[i]; +// int mata_row_len = mata_crows[i + 1] - mata_crows[i]; +// int matb_row_len = matb_crows[i + 1] - matb_crows[i]; +// if (mata_row_len == 0) { +// // cudaMemcpyAsync(&matc_cols[row_start], &matc_cols[matb_crows[i]], matb_row_len * sizeof(long), +// // cudaMemcpyDeviceToDevice); +// for (int j = 0; j < matb_row_len; j++) { +// matc_cols[row_start + j] = matb_cols[matb_crows[i] + j]; +// } +// } else if (matb_row_len == 0) { +// // cudaMemcpyAsync(&matc_cols[row_start], &mata_cols[mata_crows[i]], mata_row_len * sizeof(long), +// // cudaMemcpyDeviceToDevice); +// for (int j = 0; j < mata_row_len; j++) { +// printf(""); +// // printf("i: %d j: %d row_start: %d mata_col: %ld\n", i, j, row_start, mata_cols[mata_crows[i] + j]); +// matc_cols[row_start + j] = mata_cols[mata_crows[i] + j]; +// } +// } +// } +// } + +__global__ void ReduceSum(long *matc_crows, long *mata_crows, + long *matc_cols, long *mata_cols, + int row_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < row_count; i += stride) { + int row_id = i; + int row_len = mata_crows[row_id + 1] - mata_crows[row_id]; + for (int j = 0; j < row_len; j++) { + matc_cols[matc_crows[row_id] + j] = mata_cols[mata_crows[row_id] + j]; + } + } + + // for (int i = id; i < row_count; i += stride) { + // int row_id = i; + // int row_len = matb_crows[row_id + 1] - matb_crows[row_id]; + // for (int j = 0; j < row_len; j++) { + // matc_cols[matc_crows[row_id] + j] = matb_cols[matb_crows[row_id] + j]; + // } + // } +} + +void reduce_sum_gpu(const at::Tensor& matc_crows, const at::Tensor& mata_crows, + const at::Tensor& matb_crows, const at::Tensor& matc_cols, + const at::Tensor& mata_cols, const at::Tensor& matb_cols, + const at::Tensor &mata_nnz_rows, + const at::Tensor &matb_nnz_rows) { + + int BLOCK_SIZE = 256; + // int mata_nnz_row_count = mata_nnz_rows.sizes()[0]; + // int matb_nnz_row_count = matb_nnz_rows.sizes()[0]; + // int ROW_COUNT = std::max(mata_nnz_row_count, matb_nnz_row_count); + int ROW_COUNT = matc_crows.sizes()[0] - 1; + int BLOCK_COUNT = std::ceil(ROW_COUNT / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (ROW_COUNT == 0) { + return; + } + + ReduceSum<<>>(matc_crows.data(), + mata_crows.data(), + matc_cols.data(), + mata_cols.data(), + ROW_COUNT); + + ReduceSum<<>>(matc_crows.data(), + matb_crows.data(), + matc_cols.data(), + matb_cols.data(), + ROW_COUNT); + + CHECK_ERROR("reduce sum error") +} + +__global__ void ShadowColselect(long *sampled_frontiers, long *sampled_frontiers_rowids, + long *sampled_frontiers_crows, long *sampled_frontiers_cols, + long *adj_crows, long *adj_cols, long *mask, long *ps_batch_sizes, + int batch_size, int nnz_col_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (long i = id; i < nnz_col_count; i += stride) { + long vtx = sampled_frontiers[i]; + long row_id = sampled_frontiers_rowids[i]; + + long adj_start = adj_crows[i]; + long adj_stop = adj_crows[i + 1]; + long sampled_start = sampled_frontiers_crows[row_id]; + long sampled_stop = sampled_frontiers_crows[row_id + 1]; + long batch_id = row_id / batch_size; + for (long j = adj_start; j < adj_stop; j++) { + for (long k = sampled_start; k < sampled_stop; k++) { + if (adj_cols[j] == sampled_frontiers_cols[k]) { + mask[j] = k - ps_batch_sizes[batch_id]; + break; + } + } + } + } +} + + +// shadow_colselect_gpu(sampled_frontiers, sampled_frontier_rowids, sampled_frontier_csr.crow_indices(), +// sampled_frontier_csr.col_indices(), row_select_adj.crow_indices(), colselect_mask, +// sampled_frontiers.size(0), row_select_adj._nnz()) +void shadow_colselect_gpu(const at::Tensor& sampled_frontiers, + const at::Tensor& sampled_frontiers_rowids, + const at::Tensor& sampled_frontiers_crows, + const at::Tensor& sampled_frontiers_cols, + const at::Tensor& adj_crows, + const at::Tensor& adj_cols, + const at::Tensor& colselect_mask, + const at::Tensor& ps_batch_sizes, + int batch_size, int nnz_col_count, int nnz_count) { + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(nnz_col_count / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (nnz_col_count == 0) { + return; + } + + ShadowColselect<<>>(sampled_frontiers.data(), + sampled_frontiers_rowids.data(), + sampled_frontiers_crows.data(), + sampled_frontiers_cols.data(), + adj_crows.data(), + adj_cols.data(), + // colselect_mask.data(), + colselect_mask.data(), + ps_batch_sizes.data(), + batch_size, + nnz_col_count); + + fflush(stdout); + CHECK_ERROR("shadow colselect error") +} + + +__global__ void RowSelectCsrDupes(long *nnz_cols, long *adj_crows, long *adj_cols, long *adj_vals, + long *rowselect_crows, long *rowselect_cols, long *rowselect_vals, + int nnz_col_count) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (long i = id; i < nnz_col_count; i += stride) { + long vtx = nnz_cols[i]; + long degree = adj_crows[vtx + 1] - adj_crows[vtx]; + // might be able to use memset here + for (long j = 0; j < degree; j++) { + // mask[row_offsets[vtx] + j] = true; + rowselect_cols[rowselect_crows[i] + j] = adj_cols[adj_crows[vtx] + j]; + // rowselect_vals[rowselect_crows[i] + j] = adj_crows[vtx] + j; + rowselect_vals[rowselect_crows[i] + j] = adj_vals[adj_crows[vtx] + j]; + } + } +} +void rowselect_csr_dupes_gpu(const at::Tensor& sampled_frontiers, const at::Tensor& adj_crows, + const at::Tensor& adj_cols, const at::Tensor& adj_vals, + const at::Tensor& rowselect_crows, const at::Tensor& rowselect_cols, + const at::Tensor& rowselect_vals, int nnz_col_count, int nnz_count) { + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(nnz_col_count / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (nnz_col_count == 0 || nnz_count == 0) { + return; + } + + RowSelectCsrDupes<<>>(sampled_frontiers.data(), + adj_crows.data(), + adj_cols.data(), + adj_vals.data(), + rowselect_crows.data(), + rowselect_cols.data(), + rowselect_vals.data(), + nnz_col_count); + + CHECK_ERROR("rowselect_csr_dupes") } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("sparse_coo_tensor_gpu", &sparse_coo_tensor_gpu, "Sparse Tensor GPU-only constructor"); + m.def("sparse_coo_tensor_gpu", &sparse_coo_tensor_gpu, "Sparse COO Tensor GPU-only constructor"); + m.def("sparse_csr_tensor_gpu", &sparse_csr_tensor_gpu, "Sparse CSR Tensor GPU-only constructor"); m.def("spmm_gpu", &spmm_gpu, "SpMM wrapper for cusparse"); + m.def("spgemm_gpu", &spgemm_gpu, "SpGEMM wrapper for cusparse"); + m.def("downsample_gpu", &downsample_gpu, "Downsampling for LADIES sampling algorithm"); + m.def("compute_darts_gpu", &compute_darts_gpu, "Compute dart values for LADIES sampling algorithm"); + m.def("throw_darts_gpu", &throw_darts_gpu, "Throw darts in LADIES sampling algorithm"); + m.def("compute_darts_select_gpu", &compute_darts_select_gpu, "Compute dart values for LADIES alg selection"); + m.def("throw_darts_select_gpu", &throw_darts_select_gpu, "Throw darts for LADIES alg selection"); + m.def("compute_darts1d_gpu", &compute_darts1d_gpu, "Compute 1D dart values for LADIES sampling algorithm"); + m.def("throw_darts1d_gpu", &throw_darts1d_gpu, "Throw 1D darts in LADIES sampling algorithm"); + m.def("normalize_gpu", &normalize_gpu, "Normalize values in an array based on a second array"); + m.def("normalize_csr_gpu", &normalize_csr_gpu, "Normalize rows in a CSR matrix"); + m.def("sum_csrd_gpu", &sum_csrd_gpu, "Sum rows in a CSR matrix (double values)"); + m.def("sum_csri_gpu", &sum_csri_gpu, "Sum rows in a CSR matrix (int values)"); + m.def("shift_rowselect_gpu", &shift_rowselect_gpu, "Shift row selection output matrix col values"); + m.def("shift_colselect_gpu", &shift_colselect_gpu, "Shift col selection matrix row values"); + // m.def("scatterd_add_gpu", &scatterd_add_gpu, "Implementation of scatter_add_ for doubles"); + m.def("scatteri_add_gpu", &scatteri_add_gpu, "Implementation of scatter_add_ for ints"); + m.def("rowselect_coo_gpu", &rowselect_coo_gpu, "COO row selection for sparsity-aware spgemm"); + m.def("coogeam_gpu", &coogeam_gpu, "csrgeam2 wrapper for cusparse"); + m.def("rowselect_csr_gpu", &rowselect_csr_gpu, "CSR row selection for sparsity-aware spgemm"); + m.def("normalize_batch_gpu", &normalize_batch_gpu, "Normalize SpMM output"); + m.def("sort_dst_proc_gpu", &sort_dst_proc_gpu, "Sort vertices by destination process"); + m.def("nsparse_spgemm", &nsparse_spgemm, "CUDA SpGEMM from NSparse"); + m.def("rearrange_rows_gpu", &rearrange_rows_gpu, "Rearrange rows in in place of local spgemm"); + m.def("rearrangel_rows_gpu", &rearrangel_rows_gpu, "Rearrange rows in in place of local spgemm 64-bit"); + m.def("reduce_sum_gpu", &reduce_sum_gpu, "Sum for SAGE sa-spgemm reduction"); + m.def("shadow_colselect_gpu", &shadow_colselect_gpu, "ShaDow column selection"); + m.def("rowselect_csr_dupes_gpu", &rowselect_csr_dupes_gpu, "CSR row selection with duplicates"); } diff --git a/tests/test_ladies.py b/tests/test_ladies.py new file mode 100644 index 0000000..f98cfc0 --- /dev/null +++ b/tests/test_ladies.py @@ -0,0 +1,126 @@ +import argparse +import examples.gcn_15d +# import LADIES.pytorch_ladies + +import math +import numpy as np +import torch +import torch.distributed as dist + +eps = 1e-4 + +def main(args): + """ + All vertices in new frontier exist in aggregated neighborhood of old frontier + """ + + # Compute aggregated neighborhood + current_frontier_all, next_frontier_all, adj_matrices_all, g_loc = examples.gcn_15d.main(args) + torch.set_printoptions(edgeitems=25) + print(f"g_loc: {g_loc}") + + print("") + print("Beginning LADIES sampler test") + size = dist.get_world_size() + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + for i in range(rank_n_bulkmb): + print(f"Testing minibatch {i}") + current_frontier = current_frontier_all[i] # current_frontier for minibatch i + # next_frontier = next_frontier_all[i] # next_frontier for minibatch i + # current_frontier = current_frontier_all[i].unique(sorted=True) # current_frontier for minibatch i + next_frontier = next_frontier_all[i].unique(sorted=True) # next_frontier for minibatch i + adj_matrices = adj_matrices_all[i] # adj_matrices for minibatch i + + agg_neighbors = [] + for v in current_frontier: + vtx = v.item() + neighbors_mask = g_loc._indices()[0,:] == vtx + neighbors = (g_loc._indices()[1,:])[neighbors_mask] + agg_neighbors += neighbors.tolist() + + agg_neighbors = torch.IntTensor(agg_neighbors) + agg_neighbors = agg_neighbors.unique() + + # Verify each vertex in next_frontier exists in aggregated neighorhood + vertices_match = True + for v in next_frontier: + vtx = v.item() + if vtx not in agg_neighbors: + vertices_match = False + print(f"vtx {vtx} not in aggregated neighborhood") + + if vertices_match: + print("next_frontier within aggregated neighorhood") + + """ + All edges between the two frontiers of vertices exist in the sample + """ + adj_matrix = adj_matrices[0] # assume only one layer for now + nnz_count = 0 + edges_match = True + print(f"i: {i} adj_matrix: {adj_matrix}", flush=True) + for u in current_frontier: + for v in next_frontier: + u_vtx = u.item() + v_vtx = v.item() + + neighbors_mask = g_loc._indices()[0,:] == u_vtx + neighbors = (g_loc._indices()[1,:])[neighbors_mask] + neighbors_nnz = (g_loc._values())[neighbors_mask] + + if v_vtx in neighbors: + v_vtx_idx = (neighbors == v_vtx).nonzero().item() + if math.fabs(adj_matrix._values()[nnz_count] - neighbors_nnz[v_vtx_idx]) > eps: + print(f"nnz differs u: {u_vtx} v: {v_vtx} sample_nnz: {adj_matrix._values()[nnz_count]} g_loc_nnz: {neighbors_nnz[v_vtx_idx]}") + edges_match = False + nnz_count += 1 + + if edges_match: + print("all edges between two frontiers exist in sampled adj matrix") + + """ + No other edges exist in sample (i.e. edges with an endpoint outside the two frontiers of vertices) + """ + if nnz_count != adj_matrix._nnz(): + print("there exist other edges in sampled adj matrix") + else: + print("no other edges exist in sampled adj matrix") + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='test LADIES sampling code') + parser.add_argument("--dataset", type=str, default="Cora", + help="dataset to train") + parser.add_argument("--n-layers", type=int, default=1, + help="number of hidden gcn layers") + parser.add_argument("--batch-size", type=int, default=512, + help="number of vertices in minibatch") + parser.add_argument("--samp-num", type=int, default=64, + help="number of vertices per layer of layer-wise minibatch") + parser.add_argument('--normalize', action="store_true", + help='normalize adjacency matrix') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument('--dist-backend', default='nccl', type=str, + help='distributed backend') + parser.add_argument('--replication', default=1, type=int, + help='partitioning strategy to use') + parser.add_argument('--cuda', type=int, default=0, + help='Avaiable GPU ID') + parser.add_argument('--sample-method', type=str, default='ladies', + help='Sampled Algorithms: ladies/fastgcn/full') + parser.add_argument('--batch-num', type=int, default= 10, + help='Maximum Batch Number') + parser.add_argument('--pool-num', type=int, default= 10, + help='Number of Pool') + parser.add_argument('--n-bulkmb', type=int, default=1, + help='Number of minibatches to sample in bulk') + parser.add_argument('--n-darts', type=int, default=-1, + help='Number of darts to throw per minibatch in LADIES sampling') + parser.add_argument('--timing', action="store_true", + help='whether to turn on timers') + parser.add_argument("--gpu", type=int, default=4, + help="gpus per node") + args = parser.parse_args() + print(args) + + main(args) diff --git a/tests/test_nccl.cu b/tests/test_nccl.cu new file mode 100644 index 0000000..e7d2977 --- /dev/null +++ b/tests/test_nccl.cu @@ -0,0 +1,332 @@ +#include +#include "cuda_runtime.h" +#include "nccl.h" +#include "mpi.h" +#include +#include +#include +#include + + +#define MPICHECK(cmd) do { \ + int e = cmd; \ + if( e != MPI_SUCCESS ) { \ + printf("Failed: MPI error %s:%d '%d'\n", \ + __FILE__,__LINE__, e); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +#define CUDACHECK(cmd) do { \ + cudaError_t e = cmd; \ + if( e != cudaSuccess ) { \ + printf("Failed: Cuda error %s:%d '%s'\n", \ + __FILE__,__LINE__,cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +#define NCCLCHECK(cmd) do { \ + ncclResult_t r = cmd; \ + if (r!= ncclSuccess) { \ + printf("Failed, NCCL error %s:%d '%s'\n", \ + __FILE__,__LINE__,ncclGetErrorString(r)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + + +static uint64_t getHostHash(const char* string) { + // Based on DJB2a, result = result * 33 ^ char + uint64_t result = 5381; + for (int c = 0; string[c] != '\0'; c++){ + result = ((result << 5) + result) ^ string[c]; + } + return result; +} + + +static void getHostName(char* hostname, int maxlen) { + gethostname(hostname, maxlen); + for (int i=0; i< maxlen; i++) { + if (hostname[i] == '.') { + hostname[i] = '\0'; + return; + } + } +} + + +int main(int argc, char* argv[]) +{ + int size = 1*1024*1024*1024; + + + int myRank, nRanks, localRank = 0; + + + //initializing MPI + MPICHECK(MPI_Init(&argc, &argv)); + MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank)); + MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks)); + + + //calculating localRank based on hostname which is used in selecting a GPU + uint64_t hostHashs[nRanks]; + char hostname[1024]; + getHostName(hostname, 1024); + hostHashs[myRank] = getHostHash(hostname); + MPICHECK(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, hostHashs, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD)); + for (int p=0; p +// #include +// #include "cuda_runtime.h" +// #include "nccl.h" +// #include +// #include +// +// #define CUDACHECK(cmd) do { \ +// cudaError_t e = cmd; \ +// if( e != cudaSuccess ) { \ +// printf("Failed: Cuda error %s:%d '%s'\n", \ +// __FILE__,__LINE__,cudaGetErrorString(e)); \ +// exit(EXIT_FAILURE); \ +// } \ +// } while(0) +// +// +// #define NCCLCHECK(cmd) do { \ +// ncclResult_t r = cmd; \ +// if (r!= ncclSuccess) { \ +// printf("Failed, NCCL error %s:%d '%s'\n", \ +// __FILE__,__LINE__,ncclGetErrorString(r)); \ +// exit(EXIT_FAILURE); \ +// } \ +// } while(0) +// +// int main(int argc, char** argv) { +// +// if(argc < 3) { +// std::cout << "Please specify the number of data size in thousands and number of gpus per node"; +// return 0; +// } +// +// int n; +// int ngpus; +// n = atoi(argv[1]); +// ngpus = atoi(argv[2]); +// n *= 1000; +// +// int rank, size; +// +// float** sendbuff = (float**)malloc(ngpus * sizeof(float*)); +// float** recvbuff = (float**)malloc(ngpus * sizeof(float*)); +// cudaStream_t* s = (cudaStream_t*)malloc(sizeof(cudaStream_t)*ngpus); +// +// cudaEvent_t start[ngpus]; +// cudaEvent_t stop[ngpus]; +// +// float *h_data = new float[n](); +// for (int i = 0; i < n; i++) { +// h_data[i] = 5.0f; +// } +// +// // Initialize send/receive buffers, streams, and timers +// for (int i = 0; i < ngpus; i++) { +// CUDACHECK(cudaSetDevice(i)); +// CUDACHECK(cudaMalloc(recvbuff + i, n * sizeof(float))); +// +// if (rank == 0 && i == 0) { +// CUDACHECK(cudaMalloc(sendbuff + i, n * sizeof(float))); +// CUDACHECK(cudaMemcpy(sendbuff[i], h_data, n * sizeof(float), cudaMemcpyHostToDevice)); +// } +// +// CUDACHECK(cudaMemset(recvbuff[i], 0, n * sizeof(float))); +// CUDACHECK(cudaStreamCreate(s+i)); +// +// CUDACHECK(cudaEventCreate(&start[i])); +// CUDACHECK(cudaEventCreate(&stop[i])); +// } +// +// ncclUniqueId id; +// ncclComm_t comms[ngpus]; +// +// // Generating NCCL unique ID at one process and broadcasting it to all +// if (rank == 0) { +// ncclGetUniqueId(&id); +// } +// +// // Initialize NCCL communicator. +// NCCLCHECK(ncclGroupStart()); +// for (int i = 0; i < ngpus; i++) { +// CUDACHECK(cudaSetDevice(i)); +// NCCLCHECK(ncclCommInitRank(comms+i, size*ngpus, id, rank*ngpus + i)); +// } +// NCCLCHECK(ncclGroupEnd()); +// +// for (int j = 0; j < 1; j++) { +// // Call ncclBroadcast (ncclGroup* calls make this function as one ncclBroadcast call). +// NCCLCHECK(ncclGroupStart()); +// for (int i = 0; i < ngpus; i++) { +// CUDACHECK(cudaEventRecord(start[i], s[i])); +// NCCLCHECK(ncclBroadcast((const void*)sendbuff[0], (void*)recvbuff[i], n, ncclFloat, 0, +// comms[i], s[i])); +// } +// NCCLCHECK(ncclGroupEnd()); +// } +// +// // Synchronizing on CUDA stream to complete NCCL communication. +// for (int i = 0; i < ngpus; i++) { +// CUDACHECK(cudaSetDevice(i)); +// CUDACHECK(cudaStreamSynchronize(s[i])); +// CUDACHECK(cudaEventRecord(stop[i], s[i])); +// CUDACHECK(cudaEventSynchronize(stop[i])); +// } +// +// // Collect timings and verify broadcast worked. +// float *h_recvbuff = new float[n](); +// for (int i = 0; i < ngpus; i++) { +// CUDACHECK(cudaSetDevice(i)); +// +// float time; +// CUDACHECK(cudaEventElapsedTime(&time, start[i], stop[i])); +// time /= 1000; // seconds +// +// cudaMemcpy(h_recvbuff, recvbuff[i], n * sizeof(float), cudaMemcpyDeviceToHost); +// for (int j = 0; j < n; j++) { +// if (h_recvbuff[i] != 5.0f) { +// std::cout << "bcast error" << std::endl; +// } +// } +// std::cout << "rank: " << rank << " gpu: " << i << " size: " << (n * sizeof(float)) << " time: " << time << " bw: " << ((n * sizeof(float)) / time) << std::endl; +// } +// +// // Freeing device memory +// for (int i = 0; i < ngpus; i++) { +// if (rank == 0 && i == 0) { +// CUDACHECK(cudaFree(sendbuff[i])); +// } +// CUDACHECK(cudaFree(recvbuff[i])); +// } +// +// +// // Finalizing NCCL +// for (int i=0; i 0: + # x = torch.cuda.IntTensor(5).fill_(rank + 1) + # print(f"x: {x}") + # dist.gather(x, dst=0) + # else: + # x = torch.cuda.IntTensor(5).fill_(rank + 1) + # gather_list = [] + # for i in range(size): + # gather_list.append(torch.cuda.IntTensor(5).fill_(0)) + + # print(f"before gather_list: {gather_list}") + # dist.gather(x, gather_list) + # print(f"after gather_list: {gather_list}") + + # # (sparse) allreduce test + # # buffer_size = 5 + # buffer_size = 16213041 + # # x_ind = torch.arange(buffer_size).cuda().long() + # # x_ind = torch.stack((x_ind, x_ind)) + # # x_val = torch.cuda.FloatTensor(buffer_size).fill_(1.0) + # # x = torch.sparse_coo_tensor(x_ind, x_val, size=(buffer_size, buffer_size)) + # x = torch.cuda.FloatTensor(buffer_size).fill_(3) + # recv_list = [] + # for i in range(size): + # recv_list.append(torch.cuda.FloatTensor(buffer_size)) + # recv_list2 = [] + # for i in range(size): + # recv_list2.append(torch.cuda.FloatTensor(buffer_size)) + # # dist.all_reduce(x) + # dist.all_gather(recv_list, x) + # dist.all_gather(recv_list2, x) + # # if rank == 0: + # # dist.send(x, 2) + # # dist.recv(recv_list[0], 2) + # # x = x + recv_list[0] + # # elif rank == 2: + # # dist.recv(recv_list[0], 0) + # # dist.send(x, 0) + # # x = x + recv_list[0] + + # # if rank == 1: + # # dist.send(x, 3) + # # dist.recv(recv_list[0], 3) + # # x = x + recv_list[0] + # # elif rank == 3: + # # dist.recv(recv_list[0], 1) + # # dist.send(x, 1) + # # x = x + recv_list[0] + + # # if rank == 0: + # # dist.send(x, 1) + # # dist.recv(recv_list[0], 1) + # # x = x + recv_list[0] + # # elif rank == 1: + # # dist.recv(recv_list[0], 0) + # # dist.send(x, 0) + # # x = x + recv_list[0] + + # # if rank == 2: + # # dist.send(x, 3) + # # dist.recv(recv_list[0], 3) + # # x = x + recv_list[0] + # # elif rank == 3: + # # dist.recv(recv_list[0], 2) + # # dist.send(x, 2) + # # x = x + recv_list[0] + + # start_time(start_timer) + # # dist.all_reduce(x) + # dist.all_gather(recv_list, x) + # dist.all_gather(recv_list2, x) + # # if rank == 0: + # # dist.send(x, 1) + # # dist.recv(recv_list[0], 1) + # # elif rank == 1: + # # dist.recv(recv_list[0], 0) + # # dist.send(x, 0) + # # if rank == 0: + # # dist.send(x, 2) + # # dist.recv(recv_list[0], 2) + # # x = x + recv_list[0] + # # elif rank == 2: + # # dist.recv(recv_list[0], 0) + # # dist.send(x, 0) + # # x = x + recv_list[0] + + # # if rank == 1: + # # dist.send(x, 3) + # # dist.recv(recv_list[0], 3) + # # x = x + recv_list[0] + # # elif rank == 3: + # # dist.recv(recv_list[0], 1) + # # dist.send(x, 1) + # # x = x + recv_list[0] + + # # if rank == 0: + # # dist.send(x, 1) + # # dist.recv(recv_list[0], 1) + # # x = x + recv_list[0] + # # elif rank == 1: + # # dist.recv(recv_list[0], 0) + # # dist.send(x, 0) + # # x = x + recv_list[0] + + # # if rank == 2: + # # dist.send(x, 3) + # # dist.recv(recv_list[0], 3) + # # x = x + recv_list[0] + # # elif rank == 3: + # # dist.recv(recv_list[0], 2) + # # dist.send(x, 2) + # # x = x + recv_list[0] + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 8) / 2**30 + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + + # # scatter test + # buffer_size = 5 + # x = torch.cuda.FloatTensor(buffer_size + 3).fill_(2) + # y = torch.cuda.FloatTensor(buffer_size).fill_(3) + # if rank == 0: + # scatter_list = [x, y] + # else: + # scatter_list = None + + # if rank == 0: + # z = torch.cuda.FloatTensor(buffer_size + 3) + # else: + # z = torch.cuda.FloatTensor(buffer_size) + # dist.scatterv(z, scatter_list, src=0) + + # print(f"x: {x}") + # print(f"y: {y}") + # print(f"z: {z}") + + # # batch isend vs blocking send test + # dist.barrier() + # buffer_size = 1 * 2**30 + # run_count = 2 + # torch.cuda.profiler.cudart().cudaProfilerStart() + # print(f"blocking send") + # torch.cuda.nvtx.range_push(f"nvtx-blocksend-rank{rank}") + # for r in range(run_count): + # tensors = [] + # for i in range(size - 1): + # tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i + 1)) + + # if rank == 0: + # if r == run_count - 1: + # start_time(start_timer) + + # for i in range(size - 1): + # dist.send(tensors[i], dst=i+1) + # torch.cuda.synchronize() + # recv_tensor = torch.cuda.FloatTensor(buffer_size).fill_(0) + + # if r == run_count - 1: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 4) / 2**30 + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + # else: + # recv_tensor = torch.cuda.FloatTensor(buffer_size) + # dist.recv(recv_tensor, src=0) + # dist.barrier() + # dist.barrier() + # torch.cuda.nvtx.range_pop() + + # dist.barrier() + # print(f"batch send", flush=True) + # buffer_size = 2000 * 2**20 // size + # run_count = 2 + # torch.cuda.nvtx.range_push(f"nvtx-batchsend-rank{rank}") + # for r in range(run_count): + # tensors = [] + # for i in range(size - 1): + # tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i + 1)) + + # send_ops = [] + # for i in range(size - 1): + # send_ops.append(dist.P2POp(dist.isend, tensors[i], i + 1)) + + # if rank == 0: + # if r == run_count - 1: + # start_time(start_timer) + + # reqs = dist.batch_isend_irecv(send_ops) + # for req in reqs: + # req.wait() + # recv_tensor = torch.cuda.FloatTensor(buffer_size).fill_(0) + # torch.cuda.synchronize() + + # if r == run_count - 1: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 4) / 2**30 + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + # else: + # recv_tensor = torch.cuda.FloatTensor(buffer_size) + # recv_ops = [] + # recv_ops.append(dist.P2POp(dist.irecv, tensors[i], 0)) + # reqs = dist.batch_isend_irecv(recv_ops) + # for req in reqs: + # req.wait() + # # dist.recv(recv_tensor, src=0) + # dist.barrier() + # dist.barrier() + # torch.cuda.nvtx.range_pop() + + # print(f"scatter") + # torch.cuda.nvtx.range_push(f"nvtx-scatter-rank{rank}") + # for r in range(run_count): + # if rank == 0: + # tensors = [] + # for i in range(size): + # tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i)) + # else: + # tensors = None + + # recv_tensor = torch.cuda.FloatTensor(buffer_size) + + # if rank == 0 and r == run_count - 1: + # start_time(start_timer) + + # dist.scatter(recv_tensor, tensors, src=0) + # torch.cuda.synchronize() + + # if rank == 0 and r == run_count - 1: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 4) / 2**30 + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + # dist.barrier() + # dist.barrier() + # torch.cuda.nvtx.range_pop() + + # print(f"nonblocking isends") + # torch.cuda.nvtx.range_push(f"nvtx-nonblockisend-rank{rank}") + # for r in range(run_count): + # tensors = [] + # streams = [] + # for i in range(size - 1): + # tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i + 1)) + # streams.append(torch.cuda.Stream()) + + # if rank == 0: + # if r == run_count - 1: + # start_time(start_timer) + + # reqs = [] + # for i in range(size - 1): + # with torch.cuda.stream(streams[i]): + # reqs.append(dist.isend(tensors[i], dst=i+1)) + # for req in reqs: + # with torch.cuda.stream(streams[i]): + # req.wait() + # torch.cuda.synchronize() + # recv_tensor = torch.cuda.FloatTensor(buffer_size).fill_(0) + + # if r == run_count - 1: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 4) / 2**30 + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + # else: + # recv_tensor = torch.cuda.FloatTensor(buffer_size) + # dist.recv(recv_tensor, src=0) + + # torch.cuda.nvtx.range_pop() + # torch.cuda.profiler.cudart().cudaProfilerStop() + + # print(f"alltoall", flush=True) + # buffer_size = 1 * 2**30 // size + # run_count = 2 + # for r in range(run_count): + # tensors = [] + # for i in range(size): + # tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i)) + # tensors[-1] = torch.cuda.FloatTensor(0).fill_(i) + + # recv_tensors = [] + # for i in range(size): + # if rank == size - 1: + # recv_tensors.append(torch.cuda.FloatTensor(0)) + # else: + # recv_tensors.append(torch.cuda.FloatTensor(buffer_size)) + + # if r == run_count - 1 and rank == 0: + # start_time(start_timer) + + # # dist.scatter(recv_tensor, tensors, src=0) + # dist.all_to_all(recv_tensors, tensors) + # torch.cuda.synchronize() + + # if r == run_count - 1 and rank == 0: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 4 * size) / 2**30 + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + # dist.barrier() + + # dist.barrier() + # print(f"batched isend/irecv alltoall", flush=True) + # buffer_size = 1 * 2**30 // size + # # buffer_size = ((66 // 4) * 2**20 // 2) // (size - 1) + # run_count = 3 + # for r in range(run_count): + # send_tensors = [None] * size + # recv_tensors = [None] * size + # elem_count = 0 + # for i in range(size): + # if i != rank: + # elem_count += 2 * buffer_size + # send_tensors[i] = torch.cuda.FloatTensor(buffer_size).fill_(i + 1) + # recv_tensors[i] = torch.cuda.FloatTensor(buffer_size).fill_(0) + + # ops = [] + # for i in range(size): + # if i != rank: + # ops.append(dist.P2POp(dist.isend, send_tensors[i], i)) + # ops.append(dist.P2POp(dist.irecv, recv_tensors[i], i)) + + # if r == run_count - 1: + # start_time(start_timer) + # reqs = dist.batch_isend_irecv(ops) + # for req in reqs: + # req.wait() + # torch.cuda.synchronize() + # dist.barrier() + + # if r == run_count - 1 and rank == 0: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * size * 4) / 2**30 # assumes buffer_size is per process + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + # print(f"elem_count: {elem_count}") + + # print(f"p2p cross-node send", flush=True) + # buffer_size = 10 * 2**30 + # # buffer_size = ((66 // 4) * 2**20 // 2) // (size - 1) + # if rank == 0: + # send_tensor = torch.cuda.FloatTensor(buffer_size).fill_(1.0) + # elif rank == size - 1: + # recv_tensor = torch.cuda.FloatTensor(buffer_size).fill_(1.0) + + # dist.barrier() + # if rank == size - 1: + # start_time(start_timer) + # + # if rank == 0: + # dist.send(send_tensor, dst=size-1) + # elif rank == size - 1: + # dist.recv(recv_tensor, src=0) + # torch.cuda.synchronize() + # dist.barrier() + + # if rank == size - 1: + # seconds = stop_time(start_timer, stop_timer) / 1000 + # gb_count = (buffer_size * 4) / 2**30 # assumes buffer_size is per process + # bw = gb_count / seconds + # print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + # print(f"time(ms): {seconds * 1000}") + + print(f"allreduce", flush=True) + # buffer_size = 10 * 2**30 + buffer_size = 64 * 3803 + # send_tensor = torch.cuda.FloatTensor(64, 3803).fill_(1.0) + send_tensor = torch.cuda.FloatTensor(64, 27).fill_(1.0) + for i in range(59): + send_tensor = torch.cat((send_tensor, torch.cuda.FloatTensor(64, 64).fill_(1.0)), dim=1) + + print(f"send_tensor.size(): {send_tensor.size()}", flush=True) + print(f"send_tensor.contiguous(): {send_tensor.is_contiguous()}", flush=True) + # warmup + dist.all_reduce(send_tensor) + torch.cuda.synchronize() + dist.barrier() + + for i in range(1): + if rank == 0: + start_time(start_timer) + dist.all_reduce(send_tensor) + torch.cuda.synchronize() + dist.barrier() + + if rank == 0: + seconds = stop_time(start_timer, stop_timer) / 1000 + gb_count = (buffer_size * 4 * size) / 2**30 # assumes buffer_size is per process + bw = gb_count / seconds + print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + print(f"time(ms): {seconds * 1000}") + dist.barrier() + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='NCCL P2P test') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument("--gpu", type=int, default=4, + help="gpus per node") + + args = parser.parse_args() + # torch.set_num_threads(1) + # test_nccl(args) + mp.set_start_method("spawn", force=True) + mp.spawn(test_nccl, args=(args,), nprocs=args.gpu, join=True) diff --git a/tests/test_nsparse.py b/tests/test_nsparse.py new file mode 100644 index 0000000..77cce2f --- /dev/null +++ b/tests/test_nsparse.py @@ -0,0 +1,28 @@ +import torch +from sparse_coo_tensor_cpp import nsparse_spgemm + +mata_rows = torch.randint(0, 512, (1, 512)).cuda().squeeze() +mata_cols = torch.randint(0, 589380, (1, 512)).cuda().squeeze() +mata_indices = torch.stack((mata_rows, mata_cols)) +mata_values = torch.rand(512).float().cuda() + +matb_rows = torch.randint(0, 589380, (1, 4694891)).cuda().squeeze() +matb_cols = torch.randint(0, 9430088, (1, 4694891)).cuda().squeeze() +matb_indices = torch.stack((matb_rows, matb_cols)) +matb_values = torch.rand(4694891).float().cuda() + +mata = torch.sparse_coo_tensor(mata_indices, mata_values, size=torch.Size([512, 589380])).coalesce() +matb = torch.sparse_coo_tensor(matb_indices, matb_values, size=torch.Size([589380, 9430088])).coalesce() + +mata_csr = mata.to_sparse_csr() +matb_csr = matb.to_sparse_csr() + +matc = nsparse_spgemm(mata_csr.crow_indices().int(), mata_csr.col_indices().int(), mata_csr.values(), matb_csr.crow_indices().int(), matb_csr.col_indices().int(), matb_csr.values(), mata_csr.size(0), mata_csr.size(1), matb_csr.size(1)) + +c_crow_indices = matc[0] +c_col_indices = matc[1] +c_values = matc[2] + +print(c_crow_indices) +print(c_col_indices) +print(c_values) diff --git a/tests/test_sage.py b/tests/test_sage.py new file mode 100644 index 0000000..4fb07ea --- /dev/null +++ b/tests/test_sage.py @@ -0,0 +1,146 @@ +import argparse +import examples.gcn_15d +# import LADIES.pytorch_ladies + +import math +import numpy as np +import torch +import torch.distributed as dist + +eps = 1e-4 + +def main(args): + """ + All vertices in new frontier exist in aggregated neighborhood of old frontier + """ + + # Compute aggregated neighborhood + frontiers, adj_matrices_all, g_loc = examples.gcn_15d.main(args) + + print(f"g_loc: {g_loc}") + print("") + print("Beginning SAGE sampler test") + for l in range(args.n_layers): + print(f"Testing layer {l}", flush=True) + + current_frontier_all = frontiers[l] + next_frontier_all = frontiers[l + 1] + torch.set_printoptions(edgeitems=25) + size = dist.get_world_size() + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) * args.batch_size + + for i in range(rank_n_bulkmb): + print(f"Testing batch {i}", flush=True) + mb_id = i // args.batch_size + current_frontier = current_frontier_all[mb_id].view(current_frontier_all[mb_id].numel(), 1) + current_frontier = current_frontier[i % args.batch_size] # current_frontier for minibatch i + next_frontier = next_frontier_all[mb_id] + next_frontier = next_frontier[i % args.batch_size].unique(sorted=True) + + adj_matrices = adj_matrices_all[l][mb_id] # adj_matrices for minibatch i + + + vtx_mask = adj_matrices._indices()[0,:] == (i % args.batch_size) + vtx_indices = adj_matrices._indices()[:, vtx_mask] + vtx_values = adj_matrices._values()[vtx_mask] + adj_matrices = [torch.sparse_coo_tensor(vtx_indices, vtx_values, size=adj_matrices.size())] + + agg_neighbors = [] + for v in current_frontier: + vtx = v.item() + neighbors_mask = g_loc._indices()[0,:] == vtx + neighbors = (g_loc._indices()[1,:])[neighbors_mask] + agg_neighbors += neighbors.tolist() + + agg_neighbors = torch.IntTensor(agg_neighbors) + agg_neighbors = agg_neighbors.unique() + + # Verify each vertex in next_frontier exists in aggregated neighorhood + vertices_match = True + for v in next_frontier: + vtx = v.item() + if vtx not in agg_neighbors and vtx not in current_frontier: + vertices_match = False + print(f"vtx {vtx} not in aggregated neighborhood") + + if vertices_match: + print("next_frontier within aggregated neighorhood") + + """ + All edges between the two frontiers of vertices exist in the sample + """ + adj_matrix = adj_matrices[0] # assume only one layer for now + nnz_count = 0 + edges_match = True + + for u in current_frontier: + for v in next_frontier: + u_vtx = u.item() + v_vtx = v.item() + + neighbors_mask = g_loc._indices()[0,:] == u_vtx + neighbors = (g_loc._indices()[1,:])[neighbors_mask] + neighbors_nnz = (g_loc._values())[neighbors_mask] + + if v_vtx in neighbors: + v_vtx_idx = (neighbors == v_vtx).nonzero().item() + if math.fabs(adj_matrix._values()[nnz_count] - neighbors_nnz[v_vtx_idx]) > eps: + print(f"nnz differs u: {u_vtx} v: {v_vtx} sample_nnz: {adj_matrix._values()[nnz_count]} g_loc_nnz: {neighbors_nnz[v_vtx_idx]}") + edges_match = False + nnz_count += 1 + + if edges_match: + print("all edges between two frontiers exist in sampled adj matrix") + + """ + No other edges exist in sample (i.e. edges with an endpoint outside the two frontiers of vertices) + """ + if nnz_count != adj_matrix._nnz(): + print(f"there exist other edges in sampled adj matrix {nnz_count} {adj_matrix._nnz()}") + # print(f"current_frontier: {current_frontier}") + # print(f"next_frontier: {next_frontier}") + # print(f"adj_matrix: {adj_matrix}") + # print(f"g_loc[current_frontier]: {g_loc._indices()[1, g_loc._indices()[0,:] == current_frontier.item()]}") + else: + print("no other edges exist in sampled adj matrix") + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='test LADIES sampling code') + parser.add_argument("--dataset", type=str, default="Cora", + help="dataset to train") + parser.add_argument("--gpu", type=int, default=4, + help="gpus per node") + parser.add_argument("--n-layers", type=int, default=1, + help="number of hidden gcn layers") + parser.add_argument("--batch-size", type=int, default=512, + help="number of vertices in minibatch") + parser.add_argument("--samp-num", type=int, default=64, + help="number of vertices per layer of layer-wise minibatch") + parser.add_argument('--normalize', action="store_true", + help='normalize adjacency matrix') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument('--dist-backend', default='nccl', type=str, + help='distributed backend') + parser.add_argument('--replication', default=1, type=int, + help='partitioning strategy to use') + parser.add_argument('--cuda', type=int, default=0, + help='Avaiable GPU ID') + parser.add_argument('--sample-method', type=str, default='sage', + help='Sampled Algorithms: ladies/fastgcn/full') + parser.add_argument('--batch-num', type=int, default= 10, + help='Maximum Batch Number') + parser.add_argument('--pool-num', type=int, default= 10, + help='Number of Pool') + parser.add_argument('--n-bulkmb', type=int, default=1, + help='Number of minibatches to sample in bulk') + parser.add_argument('--n-darts', type=int, default=-1, + help='Number of darts to throw per minibatch in LADIES sampling') + parser.add_argument('--timing', action="store_true", + help='whether to turn on timers') + parser.add_argument('--baseline', action="store_true", + help='whether to avoid col selection for baseline comparison') + args = parser.parse_args() + print(args) + + main(args) diff --git a/tests/test_spgemm.py b/tests/test_spgemm.py new file mode 100644 index 0000000..d93a50b --- /dev/null +++ b/tests/test_spgemm.py @@ -0,0 +1,219 @@ +import argparse +import examples.gcn_15d +import math +import numpy as np +import os +import os.path as osp +import socket +import torch +import torch.distributed as dist +import torch_geometric.transforms as T +from torch_geometric.datasets import Planetoid, Reddit +from torch_geometric.utils import add_remaining_self_loops + +from cagnet.samplers.ladies import dist_spgemm1D +from cagnet.partitionings import Partitioning + +def main(args): + # Initialize inputs to spgemm (i.e. adjacency matrix, batch matrix) + device = torch.device('cuda') + path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', args.dataset) + + if args.dataset == "cora" or args.dataset == "reddit": + if args.dataset == "cora": + dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures()) + elif args.dataset == "reddit": + dataset = Reddit(path, transform=T.NormalizeFeatures()) + + data = dataset[0] + data = data.to(device) + data.x.requires_grad = True + inputs = data.x.to(device) + inputs.requires_grad = True + data.y = data.y.to(device) + edge_index = data.edge_index + num_features = dataset.num_features + num_classes = dataset.num_classes + adj_matrix = edge_index + elif args.dataset == "Amazon": + print(f"Loading coo...", flush=True) + edge_index = torch.load("../../data/Amazon/processed/data.pt") + print(f"Done loading coo", flush=True) + n = 14249639 + num_features = 300 + num_classes = 24 + inputs = torch.rand(n, num_features) + data = Data() + data.y = torch.rand(n).uniform_(0, num_classes - 1).long() + data.train_mask = torch.ones(n).long() + data.test_mask = torch.ones(n).long() + adj_matrix = edge_index.t_() + data = data.to(device) + inputs.requires_grad = True + data.y = data.y.to(device) + + # Initialize distributed environment with SLURM + if "SLURM_PROCID" in os.environ.keys(): + os.environ["RANK"] = os.environ["SLURM_PROCID"] + + if "SLURM_NTASKS" in os.environ.keys(): + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + + os.environ["MASTER_ADDR"] = args.hostname + os.environ["MASTER_PORT"] = "1234" + + dist.init_process_group(backend=args.dist_backend) + rank = dist.get_rank() + size = dist.get_world_size() + print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) + + g_loc, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + g_loc = g_loc.to(device) + + print("coalescing", flush=True) + g_loc = torch.sparse_coo_tensor(g_loc, torch.ones(g_loc.size(1)).cuda(), + size=(inputs.size(0), inputs.size(0)), + requires_grad=False) + g_loc = g_loc.coalesce().t() + print("normalizing", flush=True) + g_loc = examples.gcn_15d.row_normalize(g_loc) + print("done normalizing", flush=True) + + partitioning = Partitioning.ONE5D + + row_groups, col_groups = examples.gcn_15d.get_proc_groups(rank, size, args.replication) + + rank_c = rank // args.replication + rank_col = rank % args.replication + if rank_c >= (size // args.replication): + return + + batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major + node_count = g_loc.size(0) + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") + torch.manual_seed(0) + train_nid = data.train_mask.nonzero().squeeze() + vertex_perm = torch.randperm(train_nid.size(0)) + # Generate minibatch vertices + for i in range(args.n_bulkmb): + idx = vertex_perm[(i * args.batch_size):((i + 1) * args.batch_size)] + batches[i,:] = train_nid[idx] + torch.cuda.nvtx.range_pop() + + batches_loc = batches + torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") + batches_indices_rows = torch.arange(batches_loc.size(0)).cuda() + batches_indices_rows = batches_indices_rows.repeat_interleave(batches_loc.size(1)) + batches_indices_cols = batches_loc.view(-1) + batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) + batches_values = torch.cuda.FloatTensor(batches_loc.size(1) * batches_loc.size(0)).fill_(1.0) + batches_loc = torch.sparse_coo_tensor(batches_indices, batches_values, (batches_loc.size(0), g_loc.size(1))) + g_loc = torch.pow(g_loc, 2) + + group = col_groups[0] # Assume replication factor == 1 + dist.barrier(group) + + # Run SpGEMM with only rank0, pre-partitioning + if rank == 0: + ans_indices, ans_values = dist_spgemm1D(batches_loc, g_loc, 0, 1, row_groups[0]) + + dist.barrier(group) + + # Run partitioning + g_loc, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + _, g_loc, _ = examples.gcn_15d.one5d_partition(rank, size, inputs, g_loc, data, \ + inputs, num_classes, args.replication, \ + args.normalize) + g_loc = g_loc.to(device) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t() + print("normalizing", flush=True) + g_loc = examples.gcn_15d.row_normalize(g_loc) + print("done normalizing", flush=True) + + batches_loc = examples.gcn_15d.oned_partition_mb(rank, size, batches, args.n_bulkmb) + torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") + batches_indices_rows = torch.arange(batches_loc.size(0)).cuda() + batches_indices_rows = batches_indices_rows.repeat_interleave(batches_loc.size(1)) + batches_indices_cols = batches_loc.view(-1) + batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) + batches_values = torch.cuda.FloatTensor(batches_loc.size(1) * batches_loc.size(0)).fill_(1.0) + batches_loc = torch.sparse_coo_tensor(batches_indices, batches_values, (batches_loc.size(0), g_loc.size(1))) + g_loc = torch.pow(g_loc, 2) + + # group = col_groups[0] # Assume replication factor == 1 + dist.barrier(group) + + # Run SpGEMM with only rank0, pre-partitioning + test_indices, test_values = dist_spgemm1D(batches_loc, g_loc, rank, size, group) + + dist.barrier(group) + + if rank == 0: + print(f"ans_indices: {ans_indices}") + print(f"ans_values: {ans_values}") + + print(f"rank: {rank} test_indices: {test_indices}") + print(f"rank: {rank} test_values: {test_values}") + + for i in range(1, size): + if rank == i: + nnz_count = torch.cuda.LongTensor([test_values.size(0)]) + dist.send(nnz_count, dst=0) + dist.send(test_indices, dst=0) + dist.send(test_values, dst=0) + elif rank == 0: + nnz_count = torch.cuda.LongTensor([test_values.size(0)]) + dist.recv(nnz_count, src=i) + test_indices_recv = torch.cuda.LongTensor(2, nnz_count.item()) + dist.recv(test_indices_recv, src=i) + test_values_recv = torch.cuda.FloatTensor(nnz_count.item()) + dist.recv(test_values_recv, src=i) + + chunk_size = math.ceil(float(args.n_bulkmb / size)) + test_indices_recv[0, :] += i * chunk_size + test_indices = torch.cat((test_indices, test_indices_recv), dim=1) + test_values = torch.cat((test_values, test_values_recv), dim=0) + + if rank == 0: + print(f"final test_indices: {test_indices}") + print(f"final test_values: {test_values}") + print(f"finaleq test_indices: {(test_indices == ans_indices).all()}") + print(f"finaleq test_values: {(test_values == ans_values).all()}") + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='test distributed SpGEMM function') + parser.add_argument("--dataset", type=str, default="Cora", + help="dataset to train") + parser.add_argument("--n-layers", type=int, default=1, + help="number of hidden gcn layers") + parser.add_argument("--batch-size", type=int, default=512, + help="number of vertices in minibatch") + parser.add_argument("--samp-num", type=int, default=64, + help="number of vertices per layer of layer-wise minibatch") + parser.add_argument('--normalize', action="store_true", + help='normalize adjacency matrix') + parser.add_argument('--hostname', default='127.0.0.1', type=str, + help='hostname for rank 0') + parser.add_argument('--dist-backend', default='nccl', type=str, + help='distributed backend') + parser.add_argument('--replication', default=1, type=int, + help='partitioning strategy to use') + parser.add_argument('--cuda', type=int, default=0, + help='Avaiable GPU ID') + parser.add_argument('--sample-method', type=str, default='ladies', + help='Sampled Algorithms: ladies/fastgcn/full') + parser.add_argument('--batch-num', type=int, default= 10, + help='Maximum Batch Number') + parser.add_argument('--pool-num', type=int, default= 10, + help='Number of Pool') + parser.add_argument('--n-bulkmb', type=int, default=1, + help='Number of minibatches to sample in bulk') + parser.add_argument('--n-darts', type=int, default=10, + help='Number of darts to throw per minibatch in LADIES sampling') + args = parser.parse_args() + print(args) + + main(args)