diff --git a/.gitignore b/.gitignore index 5ed031f..0d499ca 100644 --- a/.gitignore +++ b/.gitignore @@ -3,11 +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 +slurm_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/*.batch +*.qdrep +*.out +*.nsys-rep +*.ncu-rep +examples/wandb +*.png diff --git a/README.md b/README.md index db6581f..5cccbe0 100644 --- a/README.md +++ b/README.md @@ -2,89 +2,76 @@ ## 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. +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 sparse-sparse matrix multiplication. All algorithms are implemented with `torch.distributed` on GPU-equipped clusters. -For more information, please read our ACM/IEEE SC'20 paper [Reducing Communication in Graph Neural Network Training](https://arxiv.org/pdf/2005.03300.pdf). +This branch contains a minibatch GNN training pipeline that implements GraphSAGE and LADIES sampling using sparse-sparse matrix multiplication. For more information, please read our MLSys'24 paper [Distributed Matrix-Based Sampling for Graph Neural Network Training](https://arxiv.org/abs/2311.02909). **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.9.19 +- PyTorch 1.13 +- PyTorch Geometric (PyG) 2.5.3 +- CUDA 11.7 +- GCC 11.2 -On OLCF Summit, all of these dependencies can be accessed with the following +On NERSC Perlmutter, you can run the following to load PyTorch and PyG. ```bash -module load cuda # CUDA 10.1 -module load gcc # GCC 6.4.0 -module load ibm-wml-ce/1.7.0-3 # PyTorch 1.3.1, Python 3.6.10 - -# PyG and its dependencies -conda create --name gnn --clone ibm-wml-ce-1.7.0-3 -conda activate gnn -pip install --no-cache-dir torch-scatter==1.4.0 -pip install --no-cache-dir torch-sparse==0.4.3 -pip install --no-cache-dir torch-cluster==1.4.5 -pip install --no-cache-dir torch-geometric==1.3.2 -``` +module load pytorch/1.13 +``` ## Compiling This code uses C++ extensions. To compile these, run ```bash cd sparse-extension -python setup.py install +mkdir build +python setup.py install --prefix build +export PYTHONPATH="/CAGNET/sparse-extension/build/lib/python3.9/site-packages/sparse_coo_tensor_cpp-0.0.0-py3.9-linux-x86_64.egg:/":$PYTHONPATH ``` ## 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 - -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 -- `--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 -- `--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). - -## Running on OLCF Summit (example) - -To run the CAGNET 1.5D algorithm on Reddit with -- 16 processes -- 100 epochs -- 16 hidden layer activations -- 2-factor replication - -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 - -`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` +Our minibatch training pipeline, including both our SAGE and LADIES implementations, can be accessed by running +- `examples/gcn_15d.py` : 1.5D algorithm + +This file also has the following flags: + +- `--n-epochs ` : Number of epochs to run training +- `--dataset ` : Graph dataset to run training on +- `--timing` : Enable timing barriers to time phases in training +- `--n-hidden ` : Number of activations in the hidden layer +- `--replication ` : Replication factor +- `--sample-method ` : Sampling algorithm to use +- `--n-layers ` : Number of hidden gnn layers +- `--batch-size ` : Number of vertices in a minibatch +- `--samp-num ` : Dash-separated list of sample number (i.e. number of neighbors to sample for SAGE, or number of vertices in the next layer for LADIES). +- `--n-bulkmb ` : Number of minibatches to sample in bulk during one step of training. +- `--semibulk ` : Number of minibatches to column extract in bulkb during one LADIES call. +- `--replicate-graph` : Replicate the input graph on each GPU. + +## Running on NERSC Perlmutter (example) + +To run the CAGNET 1.5D algorithm on ogbn-products with +- 4 GPUs +- 30 epochs +- GraphSAGE sampling +- 1024 batch size +- 15, 10, 5 sample numbers +- 3 layers +- 48 minibatches bulk sampled per step +- 256 hidden layer activations +- Replication factor 1 for feature data +- Graph replication + +run the following command with Slurm: + +`cd examples && srun -l -n 4 --cpus-per-task 32 --ntasks-per-node 4 --gpus-per-node 4 python gcn_15d.py --dataset ogbn-products --sample-method sage --n-epochs 30 --batch-size 1024 --n-layers 3 --samp-num 15-10-5 --replication 1 --n-bulkmb 48 --replicate-graph` ## Citation To cite CAGNET, please refer to: -> Alok Tripathy, Katherine Yelick, Aydın Buluç. Reducing Communication in Graph Neural Network Training. Proceedings of the International Conference for High Performance Computing, Networking, Storage, and Analysis (SC’20), 2020. +> Alok Tripathy, Katherine Yelick, Aydın Buluç. Distributed Matrix-Based Sampling for Graph Neural Network Training. Proceedings of Machine Learning and Systems (MLSys’24), 2024. 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..79e0db9 --- /dev/null +++ b/cagnet/samplers/__init__.py @@ -0,0 +1,9 @@ +from .ladies import ladies_sampler +from .sage import sage_sampler + +__all__ = [ + 'ladies_sampler', + 'sage_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..f5adc3e --- /dev/null +++ b/cagnet/samplers/sage.py @@ -0,0 +1,333 @@ +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, frontier_sum_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)) + 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_row = mb_count * batch_size + # nnz_col = batch_size + batch_size * frontier_sizes[0] + nnz_col = batch_size * frontier_sizes[0] + # print(f"l: {i} nnz_row: {nnz_row} nnz_col: {nnz_col}", flush=True) + # 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)) + nnz_row = batch_size + nnz_col = batch_size + for j in range(i + 1): + if j < i: + nnz_row += nnz_row * frontier_sizes[j] + # nnz_col += nnz_col * frontier_sizes[j] + nnz_col = nnz_row * frontier_sizes[i] + nnz_row *= mb_count + # print(f"l: {i} nnz_row: {nnz_row} nnz_col: {nnz_col}", flush=True) + 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") + # print(f"current_frontier.size: {current_frontier.size()}", flush=True) + # print(f"before next_frontier.size: {next_frontier.size()}", flush=True) + # next_frontier = torch.cat((current_frontier, next_frontier)) + # print(f"after next_frontier.size: {next_frontier.size()}", flush=True) + # print(f"after next_frontier.nnz: {next_frontier._nnz()}", flush=True) + # 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() + # print(f"after next_frontier_nnz.size: {next_frontier_nnz.size()}", flush=True) + # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) + frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=next_frontier.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), + torch.arange(nnz_row, device=gpu), + frontier_size) + # print(f"next_frontier_rows.size: {next_frontier_rows.size()}", flush=True) + 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_row, node_count_total])) + # torch.Size([nnz * mb_count, node_count_total])) + + # next_frontier = next_frontier.coalesce() + # next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) + next_frontier_select = next_frontier._indices()[1,:].view(-1) + # 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) + + 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.remainder(nnz_col) + 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_row, nnz_col])) + torch.Size([nnz_row, (nnz_row // mb_count) + nnz_col])) + # torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) + + next_frontier_select = next_frontier._indices()[1,:].view(-1) + current_frontier_increments = torch.arange(mb_count, dtype=torch.int32, device=gpu) + current_frontier_increments *= nnz_row // mb_count + current_frontier_increments = current_frontier_increments.repeat_interleave(nnz_row // mb_count) + current_frontier_rows = current_frontier._indices()[0,:] + current_frontier_increments + current_frontier_idxs = torch.stack((current_frontier_rows, current_frontier._indices()[1,:])) + + current_frontier = sparse_coo_tensor_gpu(current_frontier_idxs, current_frontier._values(), + torch.Size([nnz_row + nnz_col * mb_count, current_frontier.size(1)])) + + next_frontier_increments = torch.arange(1, mb_count + 1, dtype=torch.int32, device=gpu) + next_frontier_increments *= nnz_row // mb_count + next_frontier_increments = next_frontier_increments.repeat_interleave(nnz_col) + next_frontier_rows = next_frontier._indices()[0,:] + next_frontier_increments + next_frontier_idxs = torch.stack((next_frontier_rows, next_frontier._indices()[1,:])) + + next_frontier = sparse_coo_tensor_gpu(next_frontier_idxs, next_frontier._values(), + torch.Size([nnz_row + nnz_col * mb_count, next_frontier.size(1)])) + # next_frontier = torch.cat((current_frontier, next_frontier)) + # print(f"before current_frontier: {current_frontier}", flush=True) + # current_frontier = current_frontier.coalesce() + # print(f"after current_frontier: {current_frontier}", flush=True) + # print(f"before next_frontier: {next_frontier}", flush=True) + # print(f"before next_frontier[==0]: {next_frontier._indices()[:,next_frontier._values() == 0]}", flush=True) + # print(f"before next_frontier[==0].size: {next_frontier._indices()[:,next_frontier._values() == 0].size()}", flush=True) + # next_frontier = next_frontier.coalesce() + # print(f"after next_frontier: {next_frontier}", flush=True) + # print(f"after next_frontier[==0]: {next_frontier._indices()[:,next_frontier._values() == 0]}", flush=True) + # print(f"after next_frontier[==0].size: {next_frontier._indices()[:,next_frontier._values() == 0].size()}", flush=True) + # next_frontier = current_frontier + next_frontier + output_indices = torch.cuda.LongTensor(2, current_frontier._nnz() + next_frontier._nnz()) + output_vals = torch.cuda.FloatTensor(current_frontier._nnz() + next_frontier._nnz()) + frontier_sum_gpu(current_frontier._indices()[0,:], + current_frontier._indices()[1,:], + current_frontier._values(), + next_frontier._indices()[0,:], + next_frontier._indices()[1,:], + next_frontier._values(), + output_indices[0,:], + output_indices[1,:], + output_vals, + current_frontier._nnz(), next_frontier._nnz(), nnz_row, nnz_col, mb_count); + next_frontier = torch.sparse_coo_tensor(output_indices, output_vals, next_frontier.size()) + # next_frontier = next_frontier.coalesce() + # print(f"after2 next_frontier: {next_frontier}", flush=True) + next_frontier_select = next_frontier._indices()[1,:].view(-1) #.clone() + + # next_frontier_missing_vtxs = next_frontier._values() == 0 + # next_frontier_select[next_frontier_missing_vtxs] = -1 + + # adj_matrices[i] = adj_matrix_sample + # frontiers[i + 1] = next_frontier_select.clone() + # next_frontier_select = next_frontier_select.view(-1) + + # print(f"current_frontier_select: {current_frontier_select}", flush=True) + # print(f"current_frontier_select.size: {current_frontier_select.size()}", flush=True) + # print(f"next_frontier_select: {next_frontier_select}", flush=True) + # print(f"next_frontier_select.size: {next_frontier_select.size()}", flush=True) + if i + 1 == n_layers: + frontiers[i + 1] = next_frontier_select.view(-1, 1) + 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/utils.py b/cagnet/samplers/utils.py new file mode 100644 index 0000000..dd695bb --- /dev/null +++ b/cagnet/samplers/utils.py @@ -0,0 +1,1638 @@ +import math +import torch +import torch.distributed as dist +import torch_sparse +from collections import defaultdict + +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, 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(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.IntTensor(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"adj_row_lens: {adj_row_lens}", flush=True) + # print(f"adj_row_lens.min: {adj_row_lens.min()}", flush=True) + # print(f"adj_row_lens.max: {adj_row_lens.max()}", flush=True) + # print(f"mata_cols: {mata_cols}", flush=True) + # print(f"mata_cols.min: {mata_cols.min()}", flush=True) + # print(f"mata_cols.min: {mata_cols.min()}", flush=True) + # matc_chunk_row_lens[mata_rows] = adj_row_lens[mata_cols] + matc_chunk_row_lens[mata_rows] = adj_row_lens[mata_cols].int() + matc_chunk_crows = torch.cuda.LongTensor(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,:], + numerator._indices()[1,:], + matc_chunk_crows, + adj_matrix.crow_indices().long(), + adj_matrix.col_indices(), + matc_chunk_cols) + # 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) + 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 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/cora_nl1_rw_bias_correct_torchviz b/examples/cora_nl1_rw_bias_correct_torchviz new file mode 100644 index 0000000..77e82e3 --- /dev/null +++ b/examples/cora_nl1_rw_bias_correct_torchviz @@ -0,0 +1,37 @@ +digraph { + graph [size="12,12"] + node [align=left fontname=monospace fontsize=10 height=0.2 ranksep=0.1 shape=box style=filled] + 139659721499248 [label=" + ()" fillcolor=darkolivegreen1] + 139659748796208 [label=NllLossBackward0] + 139659748795104 -> 139659748796208 + 139659748795104 [label=SliceBackward0] + 139659748798416 -> 139659748795104 + 139659748798416 [label=LogSoftmaxBackward0] + 139659748798368 -> 139659748798416 + 139659748798368 [label=AddBackward0] + 139659748796880 -> 139659748798368 + 139659748796880 [label=AddmmBackward0] + 139659721731184 -> 139659748796880 + 139662211497872 [label="layers.0.lin_l.bias + (7)" fillcolor=lightblue] + 139662211497872 -> 139659721731184 + 139659721731184 [label=AccumulateGrad] + 139659721375120 -> 139659748796880 + 139659721375120 [label=TBackward0] + 139659721730464 -> 139659721375120 + 139659748804144 [label="layers.0.lin_l.weight + (7, 1433)" fillcolor=lightblue] + 139659748804144 -> 139659721730464 + 139659721730464 [label=AccumulateGrad] + 139659748797504 -> 139659748798368 + 139659748797504 [label=MmBackward0] + 139659721375312 -> 139659748797504 + 139659721375312 [label=TBackward0] + 139659721731232 -> 139659721375312 + 139659748682464 [label="layers.0.lin_r.weight + (7, 1433)" fillcolor=lightblue] + 139659748682464 -> 139659721731232 + 139659721731232 [label=AccumulateGrad] + 139659748796208 -> 139659721499248 +} diff --git a/examples/cora_nl1_rw_bias_correct_torchviz.svg b/examples/cora_nl1_rw_bias_correct_torchviz.svg new file mode 100644 index 0000000..6db9d75 --- /dev/null +++ b/examples/cora_nl1_rw_bias_correct_torchviz.svg @@ -0,0 +1,189 @@ + + + + + + + + + +139659721499248 + + () + + + +139659748796208 + +NllLossBackward0 + + + +139659748796208->139659721499248 + + + + + +139659748795104 + +SliceBackward0 + + + +139659748795104->139659748796208 + + + + + +139659748798416 + +LogSoftmaxBackward0 + + + +139659748798416->139659748795104 + + + + + +139659748798368 + +AddBackward0 + + + +139659748798368->139659748798416 + + + + + +139659748796880 + +AddmmBackward0 + + + +139659748796880->139659748798368 + + + + + +139659721731184 + +AccumulateGrad + + + +139659721731184->139659748796880 + + + + + +139662211497872 + +layers.0.lin_l.bias + (7) + + + +139662211497872->139659721731184 + + + + + +139659721375120 + +TBackward0 + + + +139659721375120->139659748796880 + + + + + +139659721730464 + +AccumulateGrad + + + +139659721730464->139659721375120 + + + + + +139659748804144 + +layers.0.lin_l.weight + (7, 1433) + + + +139659748804144->139659721730464 + + + + + +139659748797504 + +MmBackward0 + + + +139659748797504->139659748798368 + + + + + +139659721375312 + +TBackward0 + + + +139659721375312->139659748797504 + + + + + +139659721731232 + +AccumulateGrad + + + +139659721731232->139659721375312 + + + + + +139659748682464 + +layers.0.lin_r.weight + (7, 1433) + + + +139659748682464->139659721731232 + + + + + diff --git a/examples/cora_nl1_rw_bias_torchviz b/examples/cora_nl1_rw_bias_torchviz new file mode 100644 index 0000000..b7b7c3b --- /dev/null +++ b/examples/cora_nl1_rw_bias_torchviz @@ -0,0 +1,37 @@ +digraph { + graph [size="12,12"] + node [align=left fontname=monospace fontsize=10 height=0.2 ranksep=0.1 shape=box style=filled] + 139673143333328 [label=" + ()" fillcolor=darkolivegreen1] + 139673143179584 [label=NllLossBackward0] + 139673143181168 -> 139673143179584 + 139673143181168 [label=SliceBackward0] + 139673143180304 -> 139673143181168 + 139673143180304 [label=LogSoftmaxBackward0] + 139673143179968 -> 139673143180304 + 139673143179968 [label=AddBackward0] + 139673143179872 -> 139673143179968 + 139673143179872 [label=AddmmBackward0] + 139673142071552 -> 139673143179872 + 139675626514048 [label="layers.0.lin_l.bias + (7)" fillcolor=lightblue] + 139675626514048 -> 139673142071552 + 139673142071552 [label=AccumulateGrad] + 139673143141424 -> 139673143179872 + 139673143141424 [label=TBackward0] + 139673142074240 -> 139673143141424 + 139673166382640 [label="layers.0.lin_l.weight + (7, 1433)" fillcolor=lightblue] + 139673166382640 -> 139673142074240 + 139673142074240 [label=AccumulateGrad] + 139673143181264 -> 139673143179968 + 139673143181264 [label=MmBackward0] + 139673143181216 -> 139673143181264 + 139673143181216 [label=TBackward0] + 139673142075200 -> 139673143181216 + 139673176199232 [label="layers.0.lin_r.weight + (7, 1433)" fillcolor=lightblue] + 139673176199232 -> 139673142075200 + 139673142075200 [label=AccumulateGrad] + 139673143179584 -> 139673143333328 +} diff --git a/examples/cora_nl1_rw_bias_torchviz.svg b/examples/cora_nl1_rw_bias_torchviz.svg new file mode 100644 index 0000000..64fda26 --- /dev/null +++ b/examples/cora_nl1_rw_bias_torchviz.svg @@ -0,0 +1,189 @@ + + + + + + + + + +139673143333328 + + () + + + +139673143179584 + +NllLossBackward0 + + + +139673143179584->139673143333328 + + + + + +139673143181168 + +SliceBackward0 + + + +139673143181168->139673143179584 + + + + + +139673143180304 + +LogSoftmaxBackward0 + + + +139673143180304->139673143181168 + + + + + +139673143179968 + +AddBackward0 + + + +139673143179968->139673143180304 + + + + + +139673143179872 + +AddmmBackward0 + + + +139673143179872->139673143179968 + + + + + +139673142071552 + +AccumulateGrad + + + +139673142071552->139673143179872 + + + + + +139675626514048 + +layers.0.lin_l.bias + (7) + + + +139675626514048->139673142071552 + + + + + +139673143141424 + +TBackward0 + + + +139673143141424->139673143179872 + + + + + +139673142074240 + +AccumulateGrad + + + +139673142074240->139673143141424 + + + + + +139673166382640 + +layers.0.lin_l.weight + (7, 1433) + + + +139673166382640->139673142074240 + + + + + +139673143181264 + +MmBackward0 + + + +139673143181264->139673143179968 + + + + + +139673143181216 + +TBackward0 + + + +139673143181216->139673143181264 + + + + + +139673142075200 + +AccumulateGrad + + + +139673142075200->139673143181216 + + + + + +139673176199232 + +layers.0.lin_r.weight + (7, 1433) + + + +139673176199232->139673142075200 + + + + + diff --git a/examples/dgl_papers.py b/examples/dgl_papers.py new file mode 100644 index 0000000..e813989 --- /dev/null +++ b/examples/dgl_papers.py @@ -0,0 +1,223 @@ +import os +import ogb +import dgl +import torch +from tqdm import tqdm +from ogb.nodeproppred import DglNodePropPredDataset +from dgl.dataloading import MultiLayerFullNeighborSampler, MultiLayerNeighborSampler +from dgl.dataloading.pytorch import NodeDataLoader +from torch.utils.data.dataloader import DataLoader +from models import SAGE +from dataset import NodeSet, NbrSampleCollater +import logging +import argparse +import time +from tensorboardX import SummaryWriter +from utils import * + + +class Trainer(object): + def __init__(self, data, args, log_path): + self.args = args + self.log_path = log_path + self.graph, self.labels, self.train_idx, self.valid_idx, self.test_idx = data + + self.fanouts_train = list(map(int, args.fanouts_train.split(','))) + self.fanouts_valid = list(map(int, args.fanouts_valid.split(','))) if args.fanouts_valid is not None \ + else self.fanouts_train + self.fanouts_test = list(map(int, args.fanouts_test.split(','))) if args.fanouts_test is not None \ + else self.fanouts_train + assert len(self.fanouts_train) == len(self.fanouts_valid) == len(self.fanouts_test) + + self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + self.sage = SAGE(128, 1024, 172, len(self.fanouts_train), torch.nn.functional.leaky_relu, 0.1).to( + self.device) + + self.batch_size = 1024 + self.num_epoch = 25 + self.loss_fn = torch.nn.CrossEntropyLoss() + self.global_iter = 0 + + self.tb_writer = SummaryWriter(f'{log_path}/tensorboard') + self.log_file = open(f'{log_path}/log.txt', mode='w', buffering=1) + write_dict(args.__dict__, self.log_file) + write_dict(self.sage.profile, self.log_file) + self.num_parameters = count_parameters(self.sage) + self.log_file.write(f'num_parameters={self.num_parameters}\n') + self.valid_loss = [] + self.valid_accu = [] + + # dataloader + self.train_collater = NbrSampleCollater( + self.graph, MultiLayerNeighborSampler(fanouts=self.fanouts_train, replace=False)) + self.train_node_set = NodeSet(self.train_idx.tolist(), self.labels[self.train_idx].tolist()) + self.train_node_loader = DataLoader(dataset=self.train_node_set, batch_size=self.batch_size, + shuffle=True, num_workers=0, pin_memory=True, + collate_fn=self.train_collater.collate, drop_last=False) + + self.valid_collater = NbrSampleCollater( + self.graph, MultiLayerNeighborSampler(fanouts=self.fanouts_valid, replace=False)) + self.valid_node_set = NodeSet(self.valid_idx.tolist(), self.labels[self.valid_idx].tolist()) + self.valid_node_loader = DataLoader(dataset=self.valid_node_set, batch_size=self.batch_size, + shuffle=False, num_workers=0, pin_memory=True, + collate_fn=self.valid_collater.collate, drop_last=False) + self.test_collater = NbrSampleCollater( + self.graph, MultiLayerNeighborSampler(fanouts=self.fanouts_test, replace=False)) + self.test_node_set = NodeSet(self.test_idx.tolist(), self.labels[self.test_idx].tolist()) + self.test_node_loader = DataLoader(dataset=self.test_node_set, batch_size=8, + shuffle=False, num_workers=0, pin_memory=True, + collate_fn=self.test_collater.collate, drop_last=False) + + def run(self): + optimizer = torch.optim.Adam(self.sage.parameters(), lr=1e-3) + lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer, factor=0.8, + patience=1000, verbose=True) + for epoch in range(self.num_epoch): + self.train(epoch, optimizer, lr_scheduler) + self.save_ckpt(epoch) + self.valid() + best_epoch = torch.topk(torch.tensor(self.valid_loss), k=1, largest=False).indices.view(-1).item() + self.sage.load_state_dict(torch.load(f'{self.log_path}/ckpt/epoch{best_epoch}.ckpt')) + self.log_file.write(f'inference with checkpoint of epoch {best_epoch}\n') + test_accu = self.test() + self.tb_writer.close() + self.log_file.close() + return self.valid_accu[best_epoch], test_accu + + def train(self, epoch, optimizer, lr_scheduler): + self.sage.train() + for n_iter, (blocks, labels) in enumerate(tqdm(self.train_node_loader, desc=f'train epoch {epoch}')): + blocks = [block.to(self.device) for block in blocks] + labels = labels.to(self.device) + batch_size = labels.shape[0] + outputs = self.sage(blocks) + pred = torch.topk(outputs, k=1).indices.view(-1) + loss = self.loss_fn(outputs, labels) + accu = self.count_corrects(pred, labels) / labels.shape[0] + optimizer.zero_grad() + loss.backward() + optimizer.step() + lr_scheduler.step(loss) + self.global_iter += 1 + loss_item = loss.item() / batch_size + self.log_file.write( + f'epoch{epoch}\titer {n_iter}\t{n_iter * 100 / len(self.train_node_loader):.4f}%\tlr={lr_scheduler._last_lr}\tloss={loss_item:.8f}\taccu={accu:.6f}\n') + self.tb_writer.add_scalar('train/lr', lr_scheduler._last_lr, self.global_iter) + self.tb_writer.add_scalar('train/loss', loss_item, self.global_iter) + self.tb_writer.add_scalar('train/accu', accu, self.global_iter) + + def valid(self): + self.sage.eval() + correct_cnt = 0 + loss_item = 0 + with torch.no_grad(): + for blocks, labels in self.valid_node_loader: + blocks = [block.to(self.device) for block in blocks] + labels = labels.to(self.device) + batch_size = labels.shape[0] + outputs = self.sage(blocks) + loss = self.loss_fn(outputs, labels) + loss_item += loss.item() / batch_size + pred = torch.topk(outputs, k=1).indices.view(-1) + correct_cnt += self.count_corrects(pred, labels) + avg_loss_amount = loss_item / len(self.valid_node_loader) + accu = correct_cnt / len(self.valid_node_set) + self.log_file.write(f'valid avg_loss={avg_loss_amount:.8f} accu={accu:.6f}\n') + self.tb_writer.add_scalar('valid/loss', avg_loss_amount, self.global_iter) + self.tb_writer.add_scalar('valid/accu', accu, self.global_iter) + self.valid_loss.append(avg_loss_amount) + self.valid_accu.append(accu) + + def test(self): + self.sage.eval() + correct_cnt = 0 + loss_amount = 0 + with torch.no_grad(): + for blocks, labels in tqdm(self.test_node_loader, desc='test'): + blocks = [block.to(self.device) for block in blocks] + labels = labels.to(self.device) + batch_size = labels.shape[0] + outputs = self.sage(blocks) + loss = self.loss_fn(outputs, labels) + loss_amount += loss.item() / batch_size + pred = torch.topk(outputs, k=1).indices.view(-1) + correct_cnt += self.count_corrects(pred, labels) + avg_loss_amount = loss_amount / len(self.test_node_loader) + accu = correct_cnt / len(self.test_node_set) + self.log_file.write( + f'test avg_loss={avg_loss_amount:.8f} accu={accu:.6f}\n') + self.tb_writer.add_scalar('test/loss', avg_loss_amount, self.global_iter) + self.tb_writer.add_scalar('test/accu', accu, self.global_iter) + return accu + + def save_ckpt(self, epoch): + torch.save(self.sage.state_dict(), f'{self.log_path}/ckpt/epoch{epoch}.ckpt') + + def load_ckpt(self, path): + self.sage.load_state_dict(torch.load(path)) + + @staticmethod + def count_corrects(pred: torch.Tensor, label: torch.Tensor) -> int: + assert pred.dim() == 1 and label.dim() == 1 and pred.shape == label.shape + return ((pred == label) + 0.0).sum().item() + + +def graphdata_preprocess(dgldataset: ogb.nodeproppred.DglNodePropPredDataset): + graph, labels = dgldataset[0] + srcs, dsts = graph.all_edges() + graph.add_edges(dsts, srcs) + labels = labels.view(-1).type(torch.int) + splitted_idx = dgldataset.get_idx_split() + train_idx, val_idx, test_idx = splitted_idx["train"], splitted_idx["valid"], splitted_idx["test"] + return [graph, labels, train_idx, val_idx, test_idx] + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--version', type=str, default='sage') + parser.add_argument('--fanouts-train', type=str, default='12,12,12') + parser.add_argument('--fanouts-valid', type=str, default=None) + parser.add_argument('--fanouts-test', type=str, default='100,100,100') + args = parser.parse_args() + + time_stamp = time.strftime("%Y-%m-%d-%H:%M:%S", time.localtime()) + valid_accus = [] + test_accus = [] + + data = graphdata_preprocess( + DglNodePropPredDataset('ogbn-papers100M', root=os.path.join(os.environ['HOME'], 'data', 'OGB')) + ) + for seed in range(10): + + os.environ['PYTHONHASHSEED'] = str(seed) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + dgl.seed(seed) + dgl.random.seed(seed) + + log_path = f'log/{args.version}-{time_stamp}-seed{seed}' + + if not os.path.exists('log'): + os.mkdir('log') + if not os.path.exists(log_path): + os.mkdir(log_path) + os.mkdir(os.path.join(log_path, 'ckpt')) + backup_code(log_path) + else: + raise ValueError(f'log path {log_path} exists') + + trainer = Trainer(data, args, log_path) + valid_accu, test_accu = trainer.run() + valid_accus.append(valid_accu) + test_accus.append(test_accu) + print(f"runned 10 times") + print(f"valid accus: {valid_accus}") + print(f"test accus: {test_accus}") + print(f"average valid accu: {np.mean(valid_accus):.4f} ± {np.std(valid_accus):.4f}") + print(f"average test accu: {np.mean(test_accus):.4f} ± {np.std(test_accus):.4f}") + print(f'numbers of params: {trainer.num_parameters}') diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py new file mode 100644 index 0000000..e4f6ebf --- /dev/null +++ b/examples/gcn_15d.py @@ -0,0 +1,1894 @@ +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.datasets import Planetoid, Reddit +from torch_geometric.nn import SAGEConv +from torch_geometric.typing import SparseTensor +# from sage_conv import SAGEConv +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 +# import wandb +# wandb.init(project="cagnet") + +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)) + + if self.n_layers > 1: + self.layers.append(SAGEConv(in_feats, n_hidden)) + for _ in range(n_layers - 2): + self.layers.append(SAGEConv(n_hidden, n_hidden)) + self.layers.append(SAGEConv(n_hidden, n_classes)) + else: + self.layers.append(SAGEConv(in_feats, n_classes)) + + 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 = 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 = F.dropout(h, p=0.5, training=self.training) + + # 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): + + # for papers100M + num_nodes = features.size(0) + value = torch.arange(graph.size(1)) + # value = torch.arange(graph.col_indices().size(0)) + + # graph_t = SparseTensor(row=graph[0], col=graph[1], + # value=value, + # is_sorted=True, + # sparse_sizes=(num_nodes, num_nodes)) + subgraph_loader = NeighborSampler(graph, node_idx=test_idx, + # subgraph_loader = NeighborSampler(graph_t, node_idx=test_idx, + sizes=[20] * self.n_layers, batch_size=1024, + shuffle=False, num_workers=32) + # non_eval_timings = copy.deepcopy(self.timings) + hs = [] + batch_count = 0 + for batch_size, n_id, adjs in subgraph_loader: + h = features[n_id].to(self.device) + # h = features[n_id] + print(f"batch_count: {batch_count} total: {len(subgraph_loader)}", flush=True) + for l, layer in enumerate(self.layers): + if self.n_layers > 1: + edge_index, _, size = adjs[l] + else: + edge_index, _, size = adjs + # edge_index_cpu, _, size = adjs[l] + # del edge_index_cpu + edge_index = edge_index.to(self.device) + # edge_index = torch.stack((edge_index.storage.row(), edge_index.storage.col())) + values = torch.cuda.FloatTensor(edge_index.size(1)).fill_(1.0) + print(f"edge_index: {edge_index} size: {size}", flush=True) + adj_batch = torch.sparse_coo_tensor(edge_index, + # torch.FloatTensor(edge_index.size(1)).fill_(1.0), + values, + size) + adj_batch = adj_batch.t_().coalesce() + + # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + h = self.layers[l](h, edge_index) + + del edge_index + del values + del adj_batch + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + hs.append(h[:batch_size].cpu()) + del h + batch_count += 1 + features = torch.cat(hs, dim=0) + return features + + # # for all other datasets + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[20], batch_size=1024, + # shuffle=False, num_workers=64) + # non_eval_timings = copy.deepcopy(self.timings) + # for l, layer in enumerate(self.layers): + # hs = [] + # print(f"l: {l}", flush=True) + # for batch_size, n_id, adj in subgraph_loader: + # edge_index_cpu, _, size = adj + # edge_index = edge_index_cpu.to(self.device) + # del edge_index_cpu + # adj_batch = torch.sparse_coo_tensor(edge_index, + # torch.cuda.FloatTensor(edge_index.size(1)).fill_(1.0), + # size) + # adj_batch = adj_batch.t().coalesce() + # h = features[n_id].to(self.device) + + # # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + # h = self.layers[l](h, edge_index) + + # del edge_index + # del adj_batch + # if l != len(self.layers) - 1: + # h = CAGF.relu(h, self.partitioning) + # # hs.append(h) # GCNConv + # h_ans = h[:batch_size] + # hs.append(h_ans.cpu()) # SAGEConv + # del h + # del h_ans + # del features + # features = torch.cat(hs, dim=0) + # for tens in hs: + # del tens + # del hs + # 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=[20], batch_size=1024, + # 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) + test_idx = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/index/test_idx.pt") + test_idx = test_idx.to(device) + print(f"train_idx: {train_idx} train_idx.dtype {train_idx.dtype}", flush=True) + print(f"test_idx: {test_idx} test_idx.dtype {test_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)) + # adj_matrix = torch.sparse_coo_tensor(adj_matrix, torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0)) + # adj_matrix = adj_matrix.to_sparse_csr() + # adj_matrix = adj_matrix.to_sparse_coo() + # adj_matrix = adj_matrix._indices() + # print(adj_matrix) + + # 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_dir.pt") + adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices_dir.pt") + # adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr_reverse.pt") + # adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices_reverse.pt") + # adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr_selfloop.pt") + # adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices_selfloop.pt") + # adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr_reverse_selfloop.pt") + # adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices_reverse_selfloop.pt") + # adj_matrix_crows = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/rowptr.pt") + # adj_matrix_cols = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/col.pt") + print(f"before 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) + adj_matrix_crows = adj_matrix_crows.int() + adj_matrix_cols = adj_matrix_cols.int() + print(f"after adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) + g_loc = torch.sparse_csr_tensor(adj_matrix_crows, adj_matrix_cols, \ + adj_matrix_vals, \ + torch.Size([inputs.size(0), inputs.size(0)])) + # adj_matrix = g_loc.clone().cpu().to_sparse_coo()._indices() + g_loc_gpu = g_loc.to(device) + del adj_matrix_crows + del adj_matrix_cols + del adj_matrix_vals + del g_loc + g_loc = g_loc_gpu + 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_gpu = g_loc.to(device) + del g_loc + g_loc = g_loc_gpu + 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) + # g_loc_crows = torch.cuda.LongTensor(g_loc_meta[1].item() + 1) + dist.recv(g_loc_crows, src=src_gpu) + torch.cuda.synchronize() + + g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) + dist.recv(g_loc_cols, src=src_gpu) + torch.cuda.synchronize() + + # 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() + + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + 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_gpu = g_loc.to(device) + del g_loc_crows + del g_loc_cols + del g_loc_values + del g_loc + g_loc = g_loc_gpu + 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() + + # 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"g_loc: {g_loc}", flush=True) + print(f"features_loc.size: {features_loc.size()}", flush=True) + print(f"before barrier", flush=True) + dist.barrier() + print(f"after barrier", flush=True) + + 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 + test_nid = test_idx + print(f"train_nid.size: {train_nid.size()}") + print(f"test_nid.size: {test_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, 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) + # lr=args.lr * math.sqrt(size)) + lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer=optimizer, factor=0.8, + patience=1000, verbose=True) + + # 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_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 = 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) + # # Collect all matrices on rank 0 + # if rank == 0: + # frontiers0_recv = [frontiers_bulk[0]] + # frontiers1_recv = [frontiers_bulk[-1]] + # adjs0_recv = [adj_matrices_bulk[0].to_sparse_coo()] + # adjs1_recv = [adj_matrices_bulk[1].to_sparse_coo()] + # for i in range(1, size): + # frontiers0_recv_tens = torch.cuda.LongTensor(frontiers_bulk[0].size()).fill_(0) + # frontiers1_recv_tens = torch.cuda.LongTensor(frontiers_bulk[-1].size()).fill_(0) + # dist.recv(frontiers0_recv_tens, src=i) + # dist.recv(frontiers1_recv_tens, src=i) + # frontiers0_recv.append(frontiers0_recv_tens) + # frontiers1_recv.append(frontiers1_recv_tens) + + # adj0_nnz_recv = torch.cuda.IntTensor(1) + # adj1_nnz_recv = torch.cuda.IntTensor(1) + # dist.recv(adj0_nnz_recv, src=i) + # dist.recv(adj1_nnz_recv, src=i) + + # adj0_recv_tens = torch.cuda.LongTensor(2, adj0_nnz_recv.item()) + # adj1_recv_tens = torch.cuda.LongTensor(2, adj1_nnz_recv.item()) + # dist.recv(adj0_recv_tens, src=i) + # dist.recv(adj1_recv_tens, src=i) + + # adj0_recv = torch.sparse_coo_tensor(adj0_recv_tens, + # torch.cuda.FloatTensor(adj0_nnz_recv.item()).fill_(1.0), + # adj_matrices_bulk[0].size()) + # adj1_recv = torch.sparse_coo_tensor(adj1_recv_tens, + # torch.cuda.FloatTensor(adj1_nnz_recv.item()).fill_(1.0), + # adj_matrices_bulk[1].size()) + # adjs0_recv.append(adj0_recv) + # adjs1_recv.append(adj1_recv) + + # frontiers_bulk[0] = torch.cat(frontiers0_recv, dim=0) + # frontiers_bulk[-1] = torch.cat(frontiers1_recv, dim=0) + # adj_matrices_bulk[0] = torch.cat(adjs0_recv, dim=0) + # adj_matrices_bulk[1] = torch.cat(adjs1_recv, dim=0) + # adj_matrices_bulk[0] = adj_matrices_bulk[0].to_sparse_csr() + # adj_matrices_bulk[1] = adj_matrices_bulk[1].to_sparse_csr() + # else: + # dist.send(frontiers_bulk[0], dst=0) + # dist.send(frontiers_bulk[-1], dst=0) + # + # adj0_nnz = torch.cuda.IntTensor(1).fill_(adj_matrices_bulk[0]._nnz()) + # adj1_nnz = torch.cuda.IntTensor(1).fill_(adj_matrices_bulk[1]._nnz()) + # dist.send(adj0_nnz, dst=0) + # dist.send(adj1_nnz, dst=0) + # dist.send(adj_matrices_bulk[0].to_sparse_coo()._indices(), dst=0) + # dist.send(adj_matrices_bulk[1].to_sparse_coo()._indices(), dst=0) + + 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") + + # adj_matrix = g_loc + # frontiers_test = [] # [i][j] --- layer i mb j + # adj_matrices_test = [] # [i][j] --- layer i mb j + # for i in range(args.n_layers + 1): + # frontiers_test.append([None] * rank_n_bulkmb) + # adj_matrices_test.append([None] * rank_n_bulkmb) + 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_select_min = i * args.batch_size * np.prod(args.samp_num, dtype=int) + # src_select_max = (i + 1) * args.batch_size * np.prod(args.samp_num, dtype=int) + nnz_col = args.batch_size + for j in range(args.n_layers): + nnz_col += nnz_col * args.samp_num[j] + src_select_min = i * nnz_col + src_select_max = (i + 1) * nnz_col + + src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) + + # print(f"i: {i} batch_vtxs: {batch_vtxs}", flush=True) + # frontiers_test[0][i] = frontiers_bulk[0][batch_select_min:batch_select_max,:] + # frontiers_test[1][i] = frontiers_bulk[-1][src_select_min:src_select_max,:] + # print(f"i: {i} frontiers_test: {frontiers_test}", flush=True) + 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)) + nnz_row = args.batch_size + nnz_col = args.batch_size + for j in range(l + 1): + if j < l: + nnz_row += nnz_row * args.samp_num[j] + nnz_col += nnz_col * args.samp_num[j] + + adj_row_select_min = i * nnz_row + adj_row_select_max = (i + 1) * nnz_row + 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) + # print(f"skip_cols: {adj_sample_skip_cols}", flush=True) + # print(f"row_lens: {adj_sample_row_lens}", flush=True) + # print(f"row_lens.size: {adj_sample_row_lens.size()}", flush=True) + # print(f"row_lens == 0: {(adj_sample_row_lens == 0).sum()}", flush=True) + # print(f"adj_sample_rows: {adj_sample_rows}", flush=True) + # print(f"row_count: {row_count}", flush=True) + 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 + # print(f"adj_sample_row_lens: {adj_sample_row_lens}", flush=True) + # print(f"adj_sample_row_lens.size: {adj_sample_row_lens.size()}", flush=True) + # print(f"(adj_sample_row_lens == 0): {(adj_sample_row_lens == 0).sum()}", flush=True) + # print(f"row_count: {row_count}", flush=True) + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, row_count, + device=adj_sample_row_lens.device), + adj_sample_row_lens) + # print(f"adj_sample_row_lens: {adj_sample_rows}", flush=True) + + # 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) + col_count = nnz_row * args.samp_num[l] + adj_sample_skip_cols_prev = adj_sample_skip_cols + adj_sample_skip_cols = torch.histc(adj_sample_cols, bins=col_count) + adj_sample_skip_cols = (adj_sample_skip_cols == 0).nonzero().squeeze(1) + if l == 0: + adj_sample_skip_cols += args.batch_size + adj_sample_skip_cols_feats = adj_sample_skip_cols + else: + inc = args.batch_size + for adj in adjs: + if adj is not None: + inc += adj._indices()[1,-1] + 1 + adj_sample_skip_cols += inc + adj_sample_skip_cols_feats = torch.cat((adj_sample_skip_cols_prev, adj_sample_skip_cols)) + # print(f"adj_sample_skip_cols_feats: {adj_sample_skip_cols_feats}", flush=True) + # print(f"adj_sample_skip_cols_feats.size: {adj_sample_skip_cols_feats.size()}", flush=True) + # 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_skip_cols = adj_sample_skip_cols_feats + # print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + # print(f"adj_sample_skip_cols.size: {adj_sample_skip_cols.size()}", flush=True) + col_count = nnz_col - 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, (nnz_row, nnz_col)) + # (args.batch_size * (args.samp_num ** l), \ + # args.batch_size * (args.samp_num ** (l + 1)))) + adj_nnzs.append(adj_sample_values.size(0)) + # adj_matrices_test[l][i] = adj_matrix_sample + # 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 = torch.cat((batch_vtxs, src_vtxs)) + 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) + # adj_sample_skip_cols += args.batch_size + # features_mask[adj_sample_skip_cols] = False + features_mask[adj_sample_skip_cols_feats] = False + features_batch = features_batch[features_mask] + + # features_batch_vtxs = features_loc[batch_vtxs] + # features_batch = torch.cat((features_batch_vtxs, features_batch)) + + # for j in range(len(adjs)): + # adjs_cols = adjs[j]._indices()[1,:] + # adjs_cols += args.batch_size + # adjs_indices = torch.stack((adjs[j]._indices()[0,:], adjs_cols)) + # adjs[j] = torch.sparse_coo_tensor(adjs_indices, adjs[j]._values(), + # size=(adjs[j].size(0), adjs[j].size(1) + args.batch_size)) + for j in reversed(range(len(adjs))): + adjs_cols = adjs[j]._indices()[1,:] + if j == len(adjs) - 1: + adjs_cols += args.batch_size + else: + adjs_cols += adjs[j + 1]._indices()[1,-1] + 1 + # adjs_cols += args.batch_size + adjs_indices = torch.stack((adjs[j]._indices()[0,:], adjs_cols)) + adjs[j] = torch.sparse_coo_tensor(adjs_indices, adjs[j]._values(), + # size=(adjs[j].size(0), features_batch.size(0))) + size=(adjs[j].size(0), adjs[j].size(1))) + else: + # adj_sample_skip_cols += args.batch_size # manual fix, change this + + # features_batch_vtxs = features_loc[batch_vtxs] + # features_batch = torch.cat((features_batch_vtxs, features_batch)) + for j in reversed(range(len(adjs))): + adjs_cols = adjs[j]._indices()[1,:] + if j == len(adjs) - 1: + adjs_cols += args.batch_size + else: + adjs_cols += adjs[j + 1]._indices()[1,-1] + 1 + # adjs_cols += args.batch_size + adjs_indices = torch.stack((adjs[j]._indices()[0,:], adjs_cols)) + adjs[j] = torch.sparse_coo_tensor(adjs_indices, adjs[j]._values(), + # size=(adjs[j].size(0), features_batch.size(0))) + size=(adjs[j].size(0), adjs[j].size(1))) + + 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) + # adj_sample_skip_cols += adjs[0].size(0) + # adj_sample_skip_cols += adjs[0]._indices()[0,-1] + 1 + + # col_start = 0 + # for j in range(len(adjs)): + # adj_idx = len(adjs) - j - 1 + # col_stop = col_start + adjs[adj_idx].size(0) + # skip_cols_mask = (adj_sample_skip_cols >= col_start) & \ + # (adj_sample_skip_cols < col_stop) + # print(f"skip_cols_mask: {skip_cols_mask}", flush=True) + # adj_sample_skip_cols[skip_cols_mask] += adjs[adj_idx].size(0) + # col_start += adjs[adj_idx].size(0) + # adj_sample_skip_cols = (src_vtxs == -1).nonzero().squeeze() + features_mask[adj_sample_skip_cols_feats] = 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[batch_vtxs]) # SAGEConv + loss = F.cross_entropy(logits[:args.batch_size], data.y[batch_vtxs]) # 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 + optimizer.step() + torch.cuda.nvtx.range_push("nvtx-optstep") + + for k 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 + + # print(f"frontiers_test: {frontiers_test}", flush=True) + # print(f"adj_matrices_test: {adj_matrices_test}", flush=True) + # print(f"adj_matrix: {adj_matrix}", flush=True) + # return frontiers_test, adj_matrices_test, adj_matrix + + 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): + model.eval() + 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 args.sample_method == "sage" and args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): + if args.sample_method == "sage" and args.dataset != "Amazon" and ("Protein" not in args.dataset): + # print(f"before evaluate() start", flush=True) + # x = input() + test_nid = test_nid.cpu() + # adj_matrix = g_loc.cpu() + adj_matrix_csr = g_loc.cpu() + adj_matrix = adj_matrix_csr.to_sparse_coo()._indices() + del adj_matrix_csr + print(f"adj_matrix: {adj_matrix}", flush=True) + print(f"inputs: {inputs}", flush=True) + print(f"adj_matrix.dtype: {adj_matrix.dtype}", flush=True) + print(f"inputs.dtype: {inputs.dtype}", flush=True) + # correct_count = model.evaluate(adj_matrix, inputs, test_nid, data.y.cpu()) + if args.dataset == "ogbn-papers100M": + # model = model.cpu() + print(f"test_nid: {test_nid}", flush=True) + print(f"test_nid.size: {test_nid.size()}", flush=True) + print(f"test_nid.unique.size: {test_nid.unique().size()}", flush=True) + out = model.evaluate(adj_matrix, inputs, test_nid, data.y) + res = out.argmax(dim=-1) == data.y.cpu()[test_nid] + print(f"out: {out}", flush=True) + print(f"res: {res}", flush=True) + print(f"train_nid: {train_nid}", flush=True) + print(f"test_nid: {test_nid}", flush=True) + acc1 = 0 + acc3 = int(res.sum()) / test_nid.size(0) + del out + del res + # model = model.to(device) + else: + # model = model.cpu() + train_nid = train_nid.cpu() + out = model.evaluate(adj_matrix, inputs, test_nid, data.y) + res = out.argmax(dim=-1) == data.y.cpu() + print(f"out: {out}", flush=True) + print(f"res: {res}", flush=True) + print(f"train_nid: {train_nid}", flush=True) + print(f"test_nid: {test_nid}", flush=True) + acc1 = int(res[train_nid].sum()) / train_nid.size(0) + acc3 = int(res[test_nid].sum()) / test_nid.size(0) + # model = model.to(device) + # acc3 = correct_count / test_nid.size(0) + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {} | Train: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc1, acc3), flush=True) + train_nid = train_nid.to(device) + test_nid = test_nid.to(device) + g_loc = g_loc.to(device).to_sparse_csr() + # wandb.log({'acc': acc3, 'time': np.sum(dur), 'loss': loss.item(), 'epoch': epoch}) + 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/gcn_2d.py b/examples/gcn_2d.py new file mode 100644 index 0000000..b8416e2 --- /dev/null +++ b/examples/gcn_2d.py @@ -0,0 +1,441 @@ +import argparse +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.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, node_count, rank, size, partitioning, device, + row_groups=None, col_groups=None, transpose_group=None): + super(GCN, self).__init__() + self.layers = nn.ModuleList() + self.node_count = node_count + self.rank = rank + self.size = size + self.row_groups = row_groups + self.col_groups = col_groups + self.transpose_group = transpose_group + self.device = device + self.partitioning = partitioning + self.n_classes = n_classes + self.timings = dict() + + proc_row = CAGF.proc_row_size(size, partitioning) + proc_col = CAGF.proc_col_size(size, partitioning) + rank_row = int(rank / proc_col) + rank_col = rank % proc_col + + self.proc_row = proc_row + self.proc_col = proc_col + self.rank_row = rank_row + self.rank_col = rank_col + + # 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=None): + h = inputs + for l, layer in enumerate(self.layers): + h = layer(self, graph, h) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + + h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + return h + +def get_proc_groups(rank, size): + proc_row = CAGF.proc_row_size(size, Partitioning.TWOD) + proc_col = CAGF.proc_col_size(size, Partitioning.TWOD) + + rank_row = int(rank / proc_col) + rank_col = rank % proc_col + + row_groups = [] + col_groups = [] + + for i in range(proc_row): + 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): + col_groups.append(dist.new_group(list(range(i, size, proc_row)))) + + 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] + + return row_groups, col_groups, transpose_group + +# 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, size): + proc_row = CAGF.proc_row_size(size, Partitioning.TWOD) + proc_col = CAGF.proc_col_size(size, Partitioning.TWOD) + + 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 + +def twod_partition(rank, size, inputs, adj_matrix, normalize, device): + node_count = inputs.size(0) + proc_row = CAGF.proc_row_size(size, Partitioning.TWOD) + proc_col = CAGF.proc_col_size(size, Partitioning.TWOD) + + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) + + 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], 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_col], normalize) + + 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, chunks_per_row, dim=0) + input_partitions = [] + for i in input_rowparts: + 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 + +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": + 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 + + # 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.TWOD + + row_groups, col_groups, transpose_group = get_proc_groups(rank, size) + + proc_row = CAGF.proc_row_size(size, partitioning) + proc_col = CAGF.proc_col_size(size, partitioning) + rank_row = int(rank / proc_col) + rank_col = rank % proc_col + + features_loc, g_loc = twod_partition(rank, size, inputs, adj_matrix, args.normalize, device) + + features_loc = features_loc.to(device) + g_loc = g_loc.coalesce().to(device) + + # create GCN model + torch.manual_seed(0) + model = GCN(num_features, + args.n_hidden, + num_classes, + args.n_layers, + inputs.size(0), + rank, + size, + partitioning, + device, + row_groups=row_groups, + col_groups=col_groups, + transpose_group=transpose_group) + + # use optimizer + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + n_per_proc = inputs.size(0) // proc_row + + rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_row] + labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_row] + + train_nid = data.train_mask.nonzero().squeeze() + test_nid = data.test_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) + loss = CAGF.cross_entropy(logits, labels_rank[rank_train_mask], train_nid.size(0), \ + num_classes, partitioning, rank, [row_groups, col_groups], \ + size, data.train_mask) + + 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) + """ + total_stop = time.time() + print(f"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='TWOD', type=str, + help='partitioning strategy to use') + args = parser.parse_args() + print(args) + + main(args) diff --git a/examples/gcn_3d.py b/examples/gcn_3d.py new file mode 100644 index 0000000..021cf95 --- /dev/null +++ b/examples/gcn_3d.py @@ -0,0 +1,533 @@ +import argparse +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.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, node_count, rank, size, partitioning, device, + row_groups=None, col_groups=None, c_groups=None, transpose_group=None): + super(GCN, self).__init__() + self.layers = nn.ModuleList() + self.node_count = node_count + self.rank = rank + self.size = size + self.row_groups = row_groups + self.col_groups = col_groups + self.c_groups = c_groups + self.transpose_group = transpose_group + self.device = device + self.partitioning = partitioning + self.n_classes = n_classes + self.timings = dict() + + proc_row = CAGF.proc_row_size(size, partitioning) + proc_col = CAGF.proc_col_size(size, partitioning) + proc_c = CAGF.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 + + self.proc_row = proc_row + self.proc_col = proc_col + self.proc_c = proc_c + self.rank_row = rank_row + self.rank_col = rank_col + self.rank_c = rank_c + + # 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=None): + h = inputs + for l, layer in enumerate(self.layers): + h = layer(self, graph, h) + + return h + +def get_proc_groups(rank, size): + proc_row = CAGF.proc_row_size(size, Partitioning.THREED) + proc_col = CAGF.proc_col_size(size, Partitioning.THREED) + proc_c = CAGF.proc_c_size(size, Partitioning.THREED) + + 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))) + + 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] + + return row_groups, col_groups, c_groups, transpose_group + +# 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 + +def twod_partition(rank, size, inputs, adj_matrix, normalize, device): + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) + + node_count = inputs.size(0) + proc_row = CAGF.proc_row_size(size, Partitioning.THREED) + proc_col = CAGF.proc_col_size(size, Partitioning.THREED) + proc_c = CAGF.proc_c_size(size, Partitioning.THREED) + + # 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], 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) + + scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], + vtx_indices[rank_col], normalize) + + # 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 + +def threed_partition_loc(rank, size, inputs, adj_matrix, height, width, device): + proc_row = CAGF.proc_row_size(size, Partitioning.THREED) + proc_col = CAGF.proc_col_size(size, Partitioning.THREED) + proc_c = CAGF.proc_c_size(size, Partitioning.THREED) + + 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 + + +# 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 = CAGF.proc_row_size(size, Partitioning.THREED) + proc_col = CAGF.proc_col_size(size, Partitioning.THREED) + + 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 + +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": + 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 + + # 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.THREED + + row_groups, col_groups, c_groups, transpose_group = get_proc_groups(rank, size) + + proc_row = CAGF.proc_row_size(size, Partitioning.THREED) + proc_col = CAGF.proc_col_size(size, Partitioning.THREED) + proc_c = CAGF.proc_c_size(size, Partitioning.THREED) + + 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 + + print(f"Before partitioning...", flush=True) + features_loc, g_loc = twod_partition(rank, size, inputs, adj_matrix, args.normalize, device) + + g_loc = g_loc.coalesce() + + features_loc, g_loc = threed_partition_loc(rank, size, features_loc, g_loc.indices(), + g_loc.size(0), g_loc.size(1), + device) + print(f"After partitioning...", flush=True) + + g_loc = g_loc.coalesce() + features_loc = features_loc.to(device) + g_loc = g_loc.to(device) + + # create GCN model + torch.manual_seed(0) + model = GCN(num_features, + args.n_hidden, + num_classes, + args.n_layers, + inputs.size(0), + rank, + size, + partitioning, + device, + row_groups=row_groups, + col_groups=col_groups, + c_groups=c_groups, + transpose_group=transpose_group) + + # use optimizer + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + n_per_proc = inputs.size(0) // proc_row + + rank_train_mask = torch.split(data.train_mask, inputs.size(0), dim=0)[rank_row * proc_c + rank_c] + labels_rank = torch.split(data.y, inputs.size(0), dim=0)[rank_row * proc_c + rank_c] + + train_nid = data.train_mask.nonzero().squeeze() + test_nid = data.test_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) + loss = CAGF.cross_entropy(logits, labels_rank[rank_train_mask], train_nid.size(0), \ + num_classes, partitioning, rank, [row_groups, col_groups, c_groups], \ + size, data.train_mask) + + 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) + """ + total_stop = time.time() + print(f"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='THREED', type=str, + help='partitioning strategy to use') + args = parser.parse_args() + print(args) + + main(args) diff --git a/examples/neighbor_sampler.py b/examples/neighbor_sampler.py new file mode 100644 index 0000000..f4cbad1 --- /dev/null +++ b/examples/neighbor_sampler.py @@ -0,0 +1,200 @@ +from typing import Callable, List, NamedTuple, Optional, Tuple, Union + +import torch +from torch import Tensor + +from torch_geometric.typing import SparseTensor + + +class EdgeIndex(NamedTuple): + edge_index: Tensor + e_id: Optional[Tensor] + size: Tuple[int, int] + + def to(self, *args, **kwargs): + edge_index = self.edge_index.to(*args, **kwargs) + e_id = self.e_id.to(*args, **kwargs) if self.e_id is not None else None + return EdgeIndex(edge_index, e_id, self.size) + + +class Adj(NamedTuple): + adj_t: SparseTensor + e_id: Optional[Tensor] + size: Tuple[int, int] + + def to(self, *args, **kwargs): + adj_t = self.adj_t.to(*args, **kwargs) + e_id = self.e_id.to(*args, **kwargs) if self.e_id is not None else None + return Adj(adj_t, e_id, self.size) + + +class NeighborSampler(torch.utils.data.DataLoader): + r"""The neighbor sampler from the `"Inductive Representation Learning on + Large Graphs" `_ paper, which allows + for mini-batch training of GNNs on large-scale graphs where full-batch + training is not feasible. + + Given a GNN with :math:`L` layers and a specific mini-batch of nodes + :obj:`node_idx` for which we want to compute embeddings, this module + iteratively samples neighbors and constructs bipartite graphs that simulate + the actual computation flow of GNNs. + + More specifically, :obj:`sizes` denotes how much neighbors we want to + sample for each node in each layer. + This module then takes in these :obj:`sizes` and iteratively samples + :obj:`sizes[l]` for each node involved in layer :obj:`l`. + In the next layer, sampling is repeated for the union of nodes that were + already encountered. + The actual computation graphs are then returned in reverse-mode, meaning + that we pass messages from a larger set of nodes to a smaller one, until we + reach the nodes for which we originally wanted to compute embeddings. + + Hence, an item returned by :class:`NeighborSampler` holds the current + :obj:`batch_size`, the IDs :obj:`n_id` of all nodes involved in the + computation, and a list of bipartite graph objects via the tuple + :obj:`(edge_index, e_id, size)`, where :obj:`edge_index` represents the + bipartite edges between source and target nodes, :obj:`e_id` denotes the + IDs of original edges in the full graph, and :obj:`size` holds the shape + of the bipartite graph. + For each bipartite graph, target nodes are also included at the beginning + of the list of source nodes so that one can easily apply skip-connections + or add self-loops. + + .. warning:: + + :class:`~torch_geometric.loader.NeighborSampler` is deprecated and will + be removed in a future release. + Use :class:`torch_geometric.loader.NeighborLoader` instead. + + .. note:: + + For an example of using :obj:`NeighborSampler`, see + `examples/reddit.py + `_ or + `examples/ogbn_products_sage.py + `_. + + Args: + edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a + :class:`torch_sparse.SparseTensor` that defines the underlying + graph connectivity/message passing flow. + :obj:`edge_index` holds the indices of a (sparse) symmetric + adjacency matrix. + If :obj:`edge_index` is of type :obj:`torch.LongTensor`, its shape + must be defined as :obj:`[2, num_edges]`, where messages from nodes + :obj:`edge_index[0]` are sent to nodes in :obj:`edge_index[1]` + (in case :obj:`flow="source_to_target"`). + If :obj:`edge_index` is of type :class:`torch_sparse.SparseTensor`, + its sparse indices :obj:`(row, col)` should relate to + :obj:`row = edge_index[1]` and :obj:`col = edge_index[0]`. + The major difference between both formats is that we need to input + the *transposed* sparse adjacency matrix. + sizes ([int]): The number of neighbors to sample for each node in each + layer. If set to :obj:`sizes[l] = -1`, all neighbors are included + in layer :obj:`l`. + node_idx (LongTensor, optional): The nodes that should be considered + for creating mini-batches. If set to :obj:`None`, all nodes will be + considered. + num_nodes (int, optional): The number of nodes in the graph. + (default: :obj:`None`) + return_e_id (bool, optional): If set to :obj:`False`, will not return + original edge indices of sampled edges. This is only useful in case + when operating on graphs without edge features to save memory. + (default: :obj:`True`) + transform (callable, optional): A function/transform that takes in + a sampled mini-batch and returns a transformed version. + (default: :obj:`None`) + **kwargs (optional): Additional arguments of + :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`, + :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`. + """ + def __init__(self, edge_index: Union[Tensor, SparseTensor], + sizes: List[int], node_idx: Optional[Tensor] = None, + num_nodes: Optional[int] = None, return_e_id: bool = True, + transform: Callable = None, **kwargs): + + edge_index = edge_index.to('cpu') + + # Remove for PyTorch Lightning: + kwargs.pop('dataset', None) + kwargs.pop('collate_fn', None) + + # Save for Pytorch Lightning < 1.6: + self.edge_index = edge_index + self.node_idx = node_idx + self.num_nodes = num_nodes + + self.sizes = sizes + self.return_e_id = return_e_id + self.transform = transform + self.is_sparse_tensor = isinstance(edge_index, SparseTensor) + self.__val__ = None + + # Obtain a *transposed* `SparseTensor` instance. + if not self.is_sparse_tensor: + if (num_nodes is None and node_idx is not None + and node_idx.dtype == torch.bool): + num_nodes = node_idx.size(0) + if (num_nodes is None and node_idx is not None + and node_idx.dtype == torch.long): + num_nodes = max(int(edge_index.max()), int(node_idx.max())) + 1 + if num_nodes is None: + num_nodes = int(edge_index.max()) + 1 + + value = torch.arange(edge_index.size(1)) if return_e_id else None + print(f"before transpose", flush=True) + x = input() + self.adj_t = SparseTensor(row=edge_index[0], col=edge_index[1], + value=value, + sparse_sizes=(num_nodes, num_nodes)).t() + print(f"after transpose", flush=True) + x = input() + else: + adj_t = edge_index + if return_e_id: + self.__val__ = adj_t.storage.value() + value = torch.arange(adj_t.nnz()) + adj_t = adj_t.set_value(value, layout='coo') + self.adj_t = adj_t + + self.adj_t.storage.rowptr() + + if node_idx is None: + node_idx = torch.arange(self.adj_t.sparse_size(0)) + elif node_idx.dtype == torch.bool: + node_idx = node_idx.nonzero(as_tuple=False).view(-1) + + super().__init__( + node_idx.view(-1).tolist(), collate_fn=self.sample, **kwargs) + + def sample(self, batch): + if not isinstance(batch, Tensor): + batch = torch.tensor(batch) + + batch_size: int = len(batch) + + adjs = [] + n_id = batch + for size in self.sizes: + adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False) + e_id = adj_t.storage.value() + size = adj_t.sparse_sizes()[::-1] + if self.__val__ is not None: + adj_t.set_value_(self.__val__[e_id], layout='coo') + + if self.is_sparse_tensor: + adjs.append(Adj(adj_t, e_id, size)) + else: + row, col, _ = adj_t.coo() + edge_index = torch.stack([col, row], dim=0) + adjs.append(EdgeIndex(edge_index, e_id, size)) + + adjs = adjs[0] if len(adjs) == 1 else adjs[::-1] + out = (batch_size, n_id, adjs) + out = self.transform(*out) if self.transform is not None else out + return out + + def __repr__(self) -> str: + return f'{self.__class__.__name__}(sizes={self.sizes})' diff --git a/examples/pyg_papers.py b/examples/pyg_papers.py new file mode 100644 index 0000000..97d7902 --- /dev/null +++ b/examples/pyg_papers.py @@ -0,0 +1,426 @@ +"""Multi-node multi-GPU example on ogbn-papers100m. + +Example way to run using srun: +srun -l -N --ntasks-per-node= \ +--container-name=cont --container-image= \ +--container-mounts=/ogb-papers100m/:/workspace/dataset +python3 path_to_script.py +""" +import gc +import os +import time +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from ogb.nodeproppred import PygNodePropPredDataset +from torch.nn.parallel import DistributedDataParallel +from torch_geometric.data import Data, Dataset +from torchmetrics import Accuracy +from torch_geometric.typing import SparseTensor + +from torch_geometric.loader import NeighborLoader, NeighborSampler +from torch_geometric.nn import GCN, GraphSAGE, SAGEConv + +def init_weights(m): + if isinstance(m, torch.nn.Linear): + gain = torch.nn.init.calculate_gain('relu') + torch.nn.init.xavier_uniform_(m.weight, gain=gain) + # m.bias.data.fill_(0.01) + +class MLP(torch.nn.Module): + def __init__(self, + input_dim, + hidden_dim, + embed_dim, + num_layers: int, + act: str = 'ReLU', + bn: bool = False, + end_up_with_fc=False, + bias=True): + super(MLP, self).__init__() + self.module_list = [] + + for i in range(num_layers): + d_in = input_dim if i == 0 else hidden_dim + d_out = embed_dim if i == num_layers - 1 else hidden_dim + self.module_list.append(torch.nn.Linear(d_in, d_out, bias=bias)) + if end_up_with_fc: + continue + if bn: + self.module_list.append(torch.nn.BatchNorm1d(d_out)) + self.module_list.append(getattr(torch.nn, act)(True)) + self.module_list = torch.nn.Sequential(*self.module_list) + + def reset_parameters(self): + for x in self.module_list: + if hasattr(x, "reset_parameters"): + x.reset_parameters() + + def forward(self, x): + return self.module_list(x) + +class SAGEResInception(torch.nn.Module): + def __init__(self, in_channels, hidden_channels, out_channels, num_layers): + conv_layer = SAGEConv + kwargs = dict(bias=False) + super().__init__() + self.num_layers = num_layers + self.convs = torch.nn.ModuleList() + self.bns = torch.nn.ModuleList() + self.res_linears = torch.nn.ModuleList() + self.hidden_channels = hidden_channels + + self.convs.append(conv_layer(in_channels, hidden_channels, **kwargs)) + self.bns.append(torch.nn.BatchNorm1d(hidden_channels)) + self.res_linears.append(torch.nn.Linear(in_channels, hidden_channels)) + for _ in range(num_layers - 2): + self.convs.append(conv_layer( + hidden_channels, hidden_channels, **kwargs)) + self.bns.append(torch.nn.BatchNorm1d(hidden_channels)) + self.res_linears.append(torch.nn.Identity()) + self.convs.append(conv_layer( + hidden_channels, hidden_channels, **kwargs)) + self.bns.append(torch.nn.BatchNorm1d(hidden_channels)) + self.res_linears.append(torch.nn.Identity()) + + self.mlp = MLP(in_channels + hidden_channels * (num_layers), + 2*out_channels, out_channels, + num_layers=2, bn=True, end_up_with_fc=True, + act='LeakyReLU') + self.reset_parameters() + + def reset_parameters(self): + for conv in self.convs: + conv.reset_parameters() + conv.apply(init_weights) + for x in self.res_linears: + if isinstance(x, torch.nn.Linear): + x.reset_parameters() + for x in self.bns: + x.reset_parameters() + self.mlp.reset_parameters() + + def forward(self, _x, adjs): + _x = _x.to(torch.float) + collect = [] + end_size = adjs[-1][-1][1] + x = F.dropout(_x, p=0.1, training=self.training) + collect.append(x[:end_size]) + for i, (edge_index, _, size) in enumerate(adjs): + x_target = x[:size[1]] + x = self.convs[i]((F.dropout(x, p=0.1, training=self.training), + F.dropout(x_target, p=0.1, + training=self.training)), edge_index) + x = self.bns[i](x) + x = F.leaky_relu(x) + x = F.dropout(x, p=0.1, training=self.training) + collect.append(x[:end_size]) + x += self.res_linears[i](x_target) + return torch.log_softmax(self.mlp(torch.cat(collect, -1)), dim=-1) + + @torch.no_grad() + def inference(self, x_all, device, subgraph_loader): + pbar = tqdm(total=x_all.size(0) * self.num_layers) + pbar.set_description('Evaluating') + + for i in range(self.num_layers): + xs = [] + for batch_size, n_id, adj in subgraph_loader: + edge_index, _, size = adj.to(device) + x = x_all[n_id].to(device) + x_target = x[:size[1]] + # x = self.convs[i]((x, x_target), edge_index) + x = self.convs[i](x, edge_index) + if i != self.num_layers - 1: + x = F.relu(x) + # xs.append(x) + xs.append(x[:batch_size].cpu()) + + pbar.update(batch_size) + + x_all = torch.cat(xs, dim=0) + + pbar.close() + + return x_all + +class SAGE(torch.nn.Module): + def __init__(self, in_channels, hidden_channels, out_channels, + num_layers=2): + super(SAGE, self).__init__() + self.num_layers = num_layers + + self.convs = torch.nn.ModuleList() + if self.num_layers > 1: + self.convs.append(SAGEConv(in_channels, hidden_channels)) + for _ in range(self.num_layers - 2): + self.convs.append(SAGEConv(hidden_channels, hidden_channels)) + self.convs.append(SAGEConv(hidden_channels, out_channels)) + else: + self.convs.append(SAGEConv(in_channels, out_channels)) + + def forward(self, x, adjs): + if self.num_layers == 1: + adjs = [adjs] + for i, (edge_index, _, size) in enumerate(adjs): + x_target = x[:size[1]] # Target nodes are always placed first. + x = self.convs[i]((x.float(), x_target.float()), edge_index) + if i != self.num_layers - 1: + x = F.relu(x) + x = F.dropout(x, p=0.5, training=self.training) + return x.log_softmax(dim=-1) + + @torch.no_grad() + def inference(self, x_all, device, subgraph_loader): + pbar = tqdm(total=x_all.size(0) * self.num_layers) + pbar.set_description('Evaluating') + + for i in range(self.num_layers): + xs = [] + for batch_size, n_id, adj in subgraph_loader: + edge_index, _, size = adj.to(device) + x = x_all[n_id].to(device) + x_target = x[:size[1]] + x = self.convs[i]((x.float(), x_target.float()), edge_index) + # x = self.convs[i](x, edge_index) + if i != self.num_layers - 1: + x = F.relu(x) + # xs.append(x) + xs.append(x[:batch_size].cpu()) + + pbar.update(batch_size) + + x_all = torch.cat(xs, dim=0) + + pbar.close() + + return x_all + +def get_num_workers() -> int: + num_workers = None + if hasattr(os, "sched_getaffinity"): + try: + num_workers = len(os.sched_getaffinity(0)) // 2 + except Exception: + pass + if num_workers is None: + num_workers = os.cpu_count() // 2 + return num_workers + + +def run(world_size, data, split_idx, model, acc, wall_clock_start): + local_id = int(os.environ['LOCAL_RANK']) % 1 + rank = torch.distributed.get_rank() + print(f"local_rank: {local_id}", flush=True) + torch.cuda.set_device(local_id) + device = torch.device(local_id) + print(f"after setting local_id: {local_id}", flush=True) + if rank == 0: + print(f'Using {nprocs} GPUs...', flush=True) + + split_idx['train'] = split_idx['train'].split( + split_idx['train'].size(0) // world_size, dim=0)[rank].clone() + # split_idx['valid'] = split_idx['valid'].split( + # split_idx['valid'].size(0) // world_size, dim=0)[rank].clone() + split_idx['test'] = split_idx['test'].split( + split_idx['test'].size(0) // world_size, dim=0)[rank].clone() + + print(f"before ddp device: {device} local_id: {local_id}", flush=True) + model = DistributedDataParallel(model.to(device), device_ids=[local_id]) + print(f"after ddp", flush=True) + optimizer = torch.optim.Adam(model.parameters(), lr=0.001, + weight_decay=5e-4) + + # kwargs = dict( + # data=data, + # batch_size=1024, + # num_workers=get_num_workers(), + # num_neighbors=[15, 10], + # ) + + print(f"before neighbor loader", flush=True) + # train_loader = NeighborLoader( + # input_nodes=split_idx['train'], + # shuffle=True, + # drop_last=True, + # **kwargs, + # ) + num_nodes = data.x.size(0) + value = torch.arange(data.edge_index.size(1)) + graph_sparsetens = SparseTensor(row=data.edge_index[0,:], col=data.edge_index[1,:], + value=value, + is_sorted=True, + sparse_sizes=(num_nodes, num_nodes)) + train_loader = NeighborSampler(graph_sparsetens, node_idx=split_idx["train"], + sizes=[15, 10], batch_size=1024, + shuffle=False, drop_last=True, num_workers=get_num_workers()) + print(f"after neighbor loader", flush=True) + # val_loader = NeighborLoader(input_nodes=split_idx['valid'], **kwargs) + # test_loader = NeighborLoader(input_nodes=split_idx['test'], **kwargs) + test_loader = NeighborSampler(graph_sparsetens, node_idx=split_idx["test"], + sizes=[20, 20], batch_size=1024, + shuffle=False, drop_last=True, num_workers=get_num_workers()) + print(f"after loaders", flush=True) + + val_steps = 1000 + warmup_steps = 100 + acc = acc.to(device) + dist.barrier() + torch.cuda.synchronize() + if rank == 0: + prep_time = round(time.perf_counter() - wall_clock_start, 2) + print("Total time before training begins (prep_time)=", prep_time, + "seconds") + print("Beginning training...") + + print(f"len(train_loader): {len(train_loader)}", flush=True) + for epoch in range(1, 6): + model.train() + # for i, batch in enumerate(train_loader): + for i, (batch_size, n_id, adjs) in enumerate(train_loader): + if i == warmup_steps: + torch.cuda.synchronize() + start = time.time() + # batch = batch.to(device) + if model.module.num_layers > 1: + adjs = [adj.to(device) for adj in adjs] + else: + adjs = adjs.to(device) + optimizer.zero_grad() + # y = batch.y[:batch.batch_size].view(-1).to(torch.long) + # out = model(batch.x, batch.edge_index)[:batch.batch_size] + y = data.y[n_id[:batch_size]].to(device).long() + out = model(data.x[n_id].to(device), adjs) + loss = F.cross_entropy(out[:batch_size,:], y) + loss.backward() + optimizer.step() + + if rank == 0 and i % 10 == 0: + print(f'Epoch: {epoch:02d}, Iteration: {i}, Loss: {loss:.4f}') + + dist.barrier() + torch.cuda.synchronize() + if rank == 0: + sec_per_iter = (time.time() - start) / (i + 1 - warmup_steps) + print(f"Avg Training Iteration Time: {sec_per_iter:.6f} s/iter") + + @torch.no_grad() + def test(loader: NeighborLoader, num_steps: Optional[int] = None): + model.eval() + # for j, batch in enumerate(loader): + for j, (batch_size, n_id, adjs) in enumerate(train_loader): + if num_steps is not None and j >= num_steps: + break + # batch = batch.to(device) + # out = model(batch.x, batch.edge_index)[:batch.batch_size] + # y = batch.y[:batch.batch_size].view(-1).to(torch.long) + if model.module.num_layers > 1: + adjs = [adj.to(device) for adj in adjs] + else: + adjs = adjs.to(device) + y = data.y[n_id[:batch_size]].to(device).long() + out = model(data.x[n_id].to(device), adjs) + acc(out, y) + acc_sum = acc.compute() + return acc_sum + + # eval_acc = test(val_loader, num_steps=val_steps) + # if rank == 0: + # print(f"Val Accuracy: {eval_acc:.4f}%", ) + + acc.reset() + dist.barrier() + + test_acc = test(test_loader) + if rank == 0: + print(f"Test Accuracy: {test_acc:.4f}%", ) + + dist.barrier() + acc.reset() + torch.cuda.synchronize() + + if rank == 0: + total_time = round(time.perf_counter() - wall_clock_start, 2) + print("Total Program Runtime (total_time) =", total_time, "seconds") + print("total_time - prep_time =", total_time - prep_time, "seconds") + + +if __name__ == '__main__': + wall_clock_start = time.perf_counter() + # Setup multi-node: + if "SLURM_PROCID" in os.environ.keys(): + os.environ["RANK"] = os.environ["SLURM_PROCID"] + os.environ["LOCAL_RANK"] = os.environ["SLURM_PROCID"] + + if "SLURM_NTASKS" in os.environ.keys(): + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + + os.environ["MASTER_ADDR"] = "nid008196" + os.environ["MASTER_PORT"] = "1234" + torch.distributed.init_process_group("nccl") + rank = dist.get_rank() + nprocs = dist.get_world_size() + assert dist.is_initialized(), "Distributed cluster not initialized" + # dataset = PygNodePropPredDataset(name='ogbn-papers100M', root="/global/u1/a/alokt/data") + + device = torch.device(f'cuda:{rank % 1}') + data = Data() + graph_crows = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/rowptr.pt") + graph_cols = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/col.pt") + graph_vals = torch.FloatTensor(graph_cols.size(0)).fill_(1.0) + num_nodes = graph_crows.size(0) - 1 + graph = torch.sparse_csr_tensor(graph_crows, graph_cols, graph_vals) + data.edge_index = graph.to_sparse_coo()._indices() + data.x = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/x.pt") + # data.x = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") + data.y = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/y.pt") + split_idx = torch.load("/global/homes/a/alokt/SALIENT_plusplus_artifact/dataset/ogbn-papers100M/split_idx.pt") + + # split_idx = dataset.get_idx_split() + print(f"split_idx['train'].size: {split_idx['train'].size()}", flush=True) + print(f"split_idx['test'].size: {split_idx['test'].size()}", flush=True) + # split_idx = dict() + # split_idx["train"] = train_idx + # split_idx["test"] = test_idx + # model = GCN(dataset.num_features, 256, 2, dataset.num_classes) + # model = GraphSAGE(dataset.num_features, 256, 2, dataset.num_classes) + model = SAGE(128, 256, 172, num_layers=2).to(rank) + # model = SAGEResInception(128, 256, 172, num_layers=3).to(rank) + acc = Accuracy(task="multiclass", num_classes=172) + # model = GCN(num_features, 256, 2, num_classes) + # acc = Accuracy(task="multiclass", num_classes=num_classes) + # data = dataset[0] + + # print(f"before casting", flush=True) + # if rank == 0: + # x = input() + # dist.barrier() + # edge_index = data.edge_index + # src = edge_index[0,:] + # dst = edge_index[1,:] + # reverse_edges = torch.stack((dst, src)) + # data.edge_index = torch.cat((edge_index, reverse_edges), dim=1) + # x = data.x.to(torch.float16) + # del data.x + # gc.collect() + # data.x = x + # print(f"after casting", flush=True) + # if rank == 0: + # x = input() + # dist.barrier() + # del reverse_edges + # del edge_index + # gc.collect() + # print(f"after collect", flush=True) + # if rank == 0: + # x = input() + # dist.barrier() + print(f"data.edge_index: {data.edge_index}", flush=True) + print(f"data.x: {data.x}", flush=True) + print(f"data.edge_index.dtype: {data.edge_index.dtype}", flush=True) + print(f"data.x.dtype: {data.x.dtype}", flush=True) + print(f"data.edge_index.size: {data.edge_index.size()}", flush=True) + data.y = data.y.reshape(-1) + run(nprocs, data, split_idx, model, acc, wall_clock_start) diff --git a/examples/sage_conv.py b/examples/sage_conv.py new file mode 100644 index 0000000..74ebfe9 --- /dev/null +++ b/examples/sage_conv.py @@ -0,0 +1,166 @@ +from typing import List, Optional, Tuple, Union + +import torch.nn.functional as F +from torch import Tensor +from torch.nn import LSTM +from torch_sparse import SparseTensor, matmul + +from torch_geometric.nn.aggr import Aggregation, MultiAggregation +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.dense.linear import Linear +from torch_geometric.typing import Adj, OptPairTensor, Size + + +class SAGEConv(MessagePassing): + r"""The GraphSAGE operator from the `"Inductive Representation Learning on + Large Graphs" `_ paper + + .. math:: + \mathbf{x}^{\prime}_i = \mathbf{W}_1 \mathbf{x}_i + \mathbf{W}_2 \cdot + \mathrm{mean}_{j \in \mathcal{N(i)}} \mathbf{x}_j + + If :obj:`project = True`, then :math:`\mathbf{x}_j` will first get + projected via + + .. math:: + \mathbf{x}_j \leftarrow \sigma ( \mathbf{W}_3 \mathbf{x}_j + + \mathbf{b}) + + as described in Eq. (3) of the paper. + + Args: + in_channels (int or tuple): Size of each input sample, or :obj:`-1` to + derive the size from the first input(s) to the forward method. + A tuple corresponds to the sizes of source and target + dimensionalities. + out_channels (int): Size of each output sample. + aggr (string or Aggregation, optional): The aggregation scheme to use. + Any aggregation of :obj:`torch_geometric.nn.aggr` can be used, + *e.g.*, :obj:`"mean"`, :obj:`"max"`, or :obj:`"lstm"`. + (default: :obj:`"mean"`) + normalize (bool, optional): If set to :obj:`True`, output features + will be :math:`\ell_2`-normalized, *i.e.*, + :math:`\frac{\mathbf{x}^{\prime}_i} + {\| \mathbf{x}^{\prime}_i \|_2}`. + (default: :obj:`False`) + root_weight (bool, optional): If set to :obj:`False`, the layer will + not add transformed root node features to the output. + (default: :obj:`True`) + project (bool, optional): If set to :obj:`True`, the layer will apply a + linear transformation followed by an activation function before + aggregation (as described in Eq. (3) of the paper). + (default: :obj:`False`) + bias (bool, optional): If set to :obj:`False`, the layer will not learn + an additive bias. (default: :obj:`True`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + + Shapes: + - **inputs:** + node features :math:`(|\mathcal{V}|, F_{in})` or + :math:`((|\mathcal{V_s}|, F_{s}), (|\mathcal{V_t}|, F_{t}))` + if bipartite, + edge indices :math:`(2, |\mathcal{E}|)` + - **outputs:** node features :math:`(|\mathcal{V}|, F_{out})` or + :math:`(|\mathcal{V_t}|, F_{out})` if bipartite + """ + def __init__( + self, + in_channels: Union[int, Tuple[int, int]], + out_channels: int, + aggr: Optional[Union[str, List[str], Aggregation]] = "mean", + normalize: bool = False, + root_weight: bool = True, + project: bool = False, + bias: bool = True, + **kwargs, + ): + self.in_channels = in_channels + self.out_channels = out_channels + self.normalize = normalize + self.root_weight = root_weight + self.project = project + + if isinstance(in_channels, int): + in_channels = (in_channels, in_channels) + + if aggr == 'lstm': + kwargs.setdefault('aggr_kwargs', {}) + kwargs['aggr_kwargs'].setdefault('in_channels', in_channels[0]) + kwargs['aggr_kwargs'].setdefault('out_channels', in_channels[0]) + + super().__init__(aggr, **kwargs) + + if self.project: + self.lin = Linear(in_channels[0], in_channels[0], bias=True) + + if self.aggr is None: + self.fuse = False # No "fused" message_and_aggregate. + self.lstm = LSTM(in_channels[0], in_channels[0], batch_first=True) + + if isinstance(self.aggr_module, MultiAggregation): + aggr_out_channels = self.aggr_module.get_out_channels( + in_channels[0]) + else: + aggr_out_channels = in_channels[0] + + self.lin_l = Linear(aggr_out_channels, out_channels, bias=bias) + if self.root_weight: + self.lin_r = Linear(in_channels[1], out_channels, bias=False) + + self.reset_parameters() + + def reset_parameters(self): + if self.project: + self.lin.reset_parameters() + self.aggr_module.reset_parameters() + self.lin_l.reset_parameters() + if self.root_weight: + self.lin_r.reset_parameters() + + def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, + size: Size = None) -> Tensor: + """""" + if isinstance(x, Tensor): + x: OptPairTensor = (x, x) + + if self.project and hasattr(self, 'lin'): + x = (self.lin(x[0]).relu(), x[1]) + + # propagate_type: (x: OptPairTensor) + # print("before propagate", flush=True) + print(f"edge_index: {edge_index} edge_index.size: {edge_index.size()}", flush=True) + print(f"size: {size}", flush=True) + # print(f"x[0]: {x[0]} x[0].size: {x[0].size()}", flush=True) + # print(f"x[0].nnz_suminds: {x[0].nonzero().squeeze().sum()}", flush=True) + out = self.propagate(edge_index, x=x, size=size) + print(f"out: {out} out.size: {out.size()}", flush=True) + print(f"lin_l.weight.sum: {self.lin_l.weight.sum()}", flush=True) + # print(f"lin_l.bias.sum: {self.lin_l.bias.sum()}", flush=True) + # print(f"out.nnz_suminds: {out.nonzero().squeeze().sum()}", flush=True) + print(f"out.sum: {out.sum()}", flush=True) + out = self.lin_l(out) + print(f"out2.sum: {out.sum()}", flush=True) + + x_r = x[1] + if self.root_weight and x_r is not None: + out = out + self.lin_r(x_r) + print(f"out3.sum: {out.sum()}", flush=True) + + if self.normalize: + out = F.normalize(out, p=2., dim=-1) + + return out + + def message(self, x_j: Tensor) -> Tensor: + return x_j + + def message_and_aggregate(self, adj_t: SparseTensor, + x: OptPairTensor) -> Tensor: + print("in messageandagg", flush=True) + adj_t = adj_t.set_value(None, layout=None) + return matmul(adj_t, x[0], reduce=self.aggr) + + def __repr__(self) -> str: + return (f'{self.__class__.__name__}({self.in_channels}, ' + f'{self.out_channels}, aggr={self.aggr})') 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 d982788..0000000 --- a/gcn_distr.py +++ /dev/null @@ -1,834 +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() - rank_train_mask = torch.split(data.train_mask.bool(), outputs.size(0), dim=0)[rank] - datay_rank = torch.split(data.y, outputs.size(0), 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 3be9279..0000000 --- a/gcn_distr_15d.py +++ /dev/null @@ -1,870 +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 - - rank_train_mask = torch.split(data.train_mask.bool(), outputs.size(0), dim=0)[rank_c] - datay_rank = torch.split(data.y, outputs.size(0), 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) - 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") - 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") - 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/sparse-extension/setup.py b/sparse-extension/setup.py index e81e5fe..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.cpp'], - 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 5edf4a8..0000000 --- a/sparse-extension/sparse_coo_tensor.cpp +++ /dev/null @@ -1,232 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include - -#include "cusparse.h" - -#include - -#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 - 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 - - cusparseDnMatDescr_t matB; - CHECK_CUSPARSE(cusparseCreateDnMat(&matB, - B.size(1), // rows - b_col, // cols - B.size(1), // 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 - 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 - &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 - &alpha, // alpha - matA, // matA - matB, // matB - &beta, // beta - matC, // matC - CUDA_R_32F, // computeType - CUSPARSE_CSRMM_ALG1, // alg - d_buffer)); // buffer - - - cudaFree(d_a_csrrows); - cudaFree(d_buffer); - - // Column-major to row-major - C.set_data(C.view({c_col, c_row})); - C.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 new file mode 100644 index 0000000..fe9b852 --- /dev/null +++ b/sparse-extension/sparse_coo_tensor.cu @@ -0,0 +1,3009 @@ +#include +#include +#include +#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; + +#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);}} + +// 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 + 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)); +} + +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]; + 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) { + + auto handle = at::cuda::getCurrentCUDASparseHandle(); + + // Impl1 -- coo2csr + csrmm2 + int nnz = A_values.size(0); + + clock_t start, stop; + + int32_t *d_a_csrrows; + 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 + 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 + + cusparseDnMatDescr_t matB; + CHECK_CUSPARSE(cusparseCreateDnMat(&matB, + 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()); + C.set_data(C.view({c_row, c_col})); + + cusparseDnMatDescr_t matC; + CHECK_CUSPARSE(cusparseCreateDnMat(&matC, + n, // rows + b_col, // cols + 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 + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_CSRMM_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 + &alpha, // alpha + matA, // matA + matB, // matB + &beta, // beta + matC, // matC + CUDA_R_32F, // computeType + CUSPARSE_CSRMM_ALG1, // alg + d_buffer)); // buffer + + + cudaFree(d_a_csrrows); + cudaFree(d_buffer); + + // 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) { + + + // 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, +__global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, long *matb_crows, + int *matb_cols, int *matc_cols, size_t nnz_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (size_t 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]; + size_t 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]); + matc_cols[dst_idx + j] = matb_cols[src_idx + j]; + // if (matc_cols[dst_idx + j] < 0) { + // printf("kernel error j: %d src_idx: %d dst_idx: %ld\n", j, src_idx, dst_idx); + // printf("kernel error dst: %d src: %d\n", matc_cols[dst_idx + j], 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_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 CurrCopy(long *curr_rows, + long *curr_cols, + float *curr_vals, + long *output_rows, + long *output_cols, + float *output_vals, + int curr_nnz, int nnz_row, int nnz_col, int mb_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + // mb_id = i / (nnz_row / mb_count) + // mb_row_start = mb_id * ((nnz_row / mb_count) + nnz_col) + // dst_idx = mb_row_start + (i % (nnz_row / mb_count) + + for (int i = id; i < curr_nnz; i += stride) { + int mb_id = i / (nnz_row / mb_count); + int mb_row_start = mb_id * ((nnz_row / mb_count) + nnz_col); + int dst_idx = mb_row_start + (i % (nnz_row / mb_count)); + output_rows[dst_idx] = curr_rows[i]; + output_cols[dst_idx] = curr_cols[i]; + output_vals[dst_idx] = curr_vals[i]; + } +} + +__global__ void NextCopy(long *next_rows, + long *next_cols, + long *next_vals, + long *output_rows, + long *output_cols, + float *output_vals, + int next_nnz, int nnz_row, int nnz_col, int mb_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + // mb_id = i / (nnz_col) + // mb_row_start = mb_id * ((nnz_row / mb_count) + nnz_col) + nnz_row + // dst_idx = mb_row_start + (i % nnz_col) + + for (int i = id; i < next_nnz; i += stride) { + int mb_id = i / nnz_col; + int mb_row_start = mb_id * ((nnz_row / mb_count) + nnz_col) + (nnz_row / mb_count); + int dst_idx = mb_row_start + (i % nnz_col); + output_rows[dst_idx] = next_rows[i]; + output_cols[dst_idx] = next_cols[i]; + output_vals[dst_idx] = (float) next_vals[i]; + } +} + +void frontier_sum_gpu(const at::Tensor& curr_rows, + const at::Tensor& curr_cols, + const at::Tensor& curr_vals, + const at::Tensor& next_rows, + const at::Tensor& next_cols, + const at::Tensor& next_vals, + const at::Tensor& output_rows, + const at::Tensor& output_cols, + const at::Tensor& output_vals, + int curr_nnz, int next_nnz, int nnz_row, int nnz_col, int mb_count) { + + 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 CURR_BLOCK_COUNT = std::ceil(curr_nnz / ((float) BLOCK_SIZE)); + CURR_BLOCK_COUNT = std::min(CURR_BLOCK_COUNT, 65535); + int NEXT_BLOCK_COUNT = std::ceil(next_nnz / ((float) BLOCK_SIZE)); + NEXT_BLOCK_COUNT = std::min(NEXT_BLOCK_COUNT, 65535); + + if (curr_nnz == 0 || next_nnz == 0) { + return; + } + + CurrCopy<<>>(curr_rows.data(), + curr_cols.data(), + curr_vals.data(), + output_rows.data(), + output_cols.data(), + output_vals.data(), + curr_nnz, nnz_row, nnz_col, mb_count); + + NextCopy<<>>(next_rows.data(), + next_cols.data(), + next_vals.data(), + output_rows.data(), + output_cols.data(), + output_vals.data(), + next_nnz, nnz_row, nnz_col, mb_count); + + CHECK_ERROR("frontier sum error") +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + 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("frontier_sum_gpu", &frontier_sum_gpu, "Summing current and next frontiers in SAGE"); +} 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_sizes = [int(2**i) for i in range(1, 15)] + # for buff_size in buffer_sizes: + # buffer_size = buff_size * 2**10 + # run_count = 2 + # for r in range(run_count): + # 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 and r == run_count - 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 r == run_count - 1 and 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"size: {buff_size} gb: {gb_count} GB time: {seconds}s bw: {bw}") + # # print(f"time(ms): {seconds * 1000}") + + print(f"cpu2gpu send", flush=True) + buffer_sizes = [int(2**i) for i in range(1, 20)] + for buff_size in buffer_sizes: + buffer_size = buff_size * 2**10 + run_count = 2 + for r in range(run_count): + tensor_buffer = torch.FloatTensor(buffer_size).fill_(1.0) + + torch.cuda.synchronize() + if r == run_count - 1: + start_time(start_timer) + tensor_buffer_dst = tensor_buffer.to(device) + if r == run_count - 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"size: {size} gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + +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() + test_nccl(args) 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..681688e --- /dev/null +++ b/tests/test_sage.py @@ -0,0 +1,170 @@ +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) + + g_loc = g_loc.to_sparse_coo() + 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 + + print(f"current_frontier: {current_frontier}", flush=True) + print(f"next_frontier: {next_frontier}", flush=True) + + 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("--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/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)