From b700bb366fbc1c6e180f5c1c9a06421b2f326408 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 10 Jan 2022 12:59:52 -0800 Subject: [PATCH 001/178] Added refactored gcn.py and gcn_conv.py --- examples/gcn.py | 359 +++++++++++++++++++++++++++++++++++++++++++ examples/gcn_conv.py | 80 ++++++++++ 2 files changed, 439 insertions(+) create mode 100644 examples/gcn.py create mode 100644 examples/gcn_conv.py diff --git a/examples/gcn.py b/examples/gcn.py new file mode 100644 index 0000000..aff888a --- /dev/null +++ b/examples/gcn.py @@ -0,0 +1,359 @@ +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.nn.functional as F +import torch.autograd.profiler as profiler +import torch_geometric +from torch_geometric.datasets import Planetoid, PPI +import torch_geometric.transforms as T +from torch_geometric.utils import add_remaining_self_loops +import torch_sparse +from gcn_conv import GCNConv +import socket + +class GCN(nn.Module): + def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, group, device): + super(GCN, self).__init__() + self.layers = nn.ModuleList() + self.rank = rank + self.size = size + self.group = group + self.device = device + self.timings = dict() + + # input layer + self.layers.append(GCNConv(in_feats, n_hidden, self.device)) + # hidden layers + for i in range(n_layers - 1): + self.layers.append(GCNConv(n_hidden, n_hidden, self.device)) + # output layer + self.layers.append(GCNConv(n_hidden, n_classes, self.device)) + + # # input layer + # self.mlp_layers.append(GCNConvMLP(in_feats, n_hidden, self.device)) + # # hidden layers + # for i in range(n_layers - 1): + # self.mlp_layers.append(GCNConvMLP(n_hidden, n_hidden, self.device)) + # # output layer + # self.mlp_layers.append(GCNConvMLP(n_hidden, n_classes, self.device)) + + def forward(self, graph, inputs, ampbyp): + h = inputs + for l, layer in enumerate(self.layers): + h = layer(self, graph, ampbyp, h) + if l != len(self.layers) - 1: + h = F.relu(h) + + h = F.log_softmax(h, 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') + graphname = "Cora" + 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 + 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)) + + 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) + + 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, + group, + device) + + # 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] + rank_test_mask = torch.split(data.test_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() + rank_test_nids = rank_test_mask.nonzero().squeeze() + + 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) + model.train() + + # forward + logits = model(g_loc, features_loc, ampbyp) + loss = F.cross_entropy(logits[rank_train_nids], labels_rank[rank_train_nids], reduction="sum") + + loss_recv = [] + for i in range(size // args.replication): + loss_recv.append(torch.cuda.FloatTensor(loss.size())) + dist.all_gather(loss_recv, loss, group) + loss_recv[rank] = loss + loss = sum(loss_recv) / train_nid.size(0) + + 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) + """ + + """ + # 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("--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("--replication", type=int, default=1, + help="replciation factor for 1.5D algorithm") + 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='gloo', 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') + args = parser.parse_args() + print(args) + + main(args) diff --git a/examples/gcn_conv.py b/examples/gcn_conv.py new file mode 100644 index 0000000..6d70180 --- /dev/null +++ b/examples/gcn_conv.py @@ -0,0 +1,80 @@ +import math +import torch +import torch.distributed as dist +import torch.nn as nn + +from sparse_coo_tensor_cpp import spmm_gpu + +def broad_func(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 outer_product(mata, matb, group): + matc = torch.mm(mata, matb) + dist.all_reduce(matc, op=dist.reduce_op.SUM, group=group) + + return matc + +class GCNConv(nn.Module): + def __init__(self, in_feats, out_feats, device): + 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) + + def forward(self, gcn, graph, ampbyp, inputs): + return GCNFunc.apply(gcn, graph, ampbyp, inputs, self.weight) + + +class GCNFunc(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, ampbyp, inputs, weight): + # inputs: H + # graph: A + # weight; W + + z = broad_func(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(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 From a1b826129980960ee9709a802d2995ef41fe5efd Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 10 Jan 2022 13:10:19 -0800 Subject: [PATCH 002/178] Removed replication factor from 1D --- examples/gcn.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/examples/gcn.py b/examples/gcn.py index aff888a..d6edda7 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -339,8 +339,6 @@ def main(args): help="Weight for L2 loss") parser.add_argument("--aggregator-type", type=str, default="gcn", help="Aggregator type: mean/gcn/pool/lstm") - parser.add_argument("--replication", type=int, default=1, - help="replciation factor for 1.5D algorithm") parser.add_argument('--world-size', default=-1, type=int, help='number of nodes for distributed training') parser.add_argument('--rank', default=-1, type=int, From e8838191b1394dfada7acb21bd04af6515964c03 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 10 Jan 2022 13:31:20 -0800 Subject: [PATCH 003/178] Turned cagnet into package --- cagnet/__init__.py | 1 + cagnet/nn/conv/__init__.py | 7 +++++++ {examples => cagnet/nn/conv}/gcn_conv.py | 0 examples/gcn.py | 4 +++- 4 files changed, 11 insertions(+), 1 deletion(-) create mode 100644 cagnet/__init__.py create mode 100644 cagnet/nn/conv/__init__.py rename {examples => cagnet/nn/conv}/gcn_conv.py (100%) diff --git a/cagnet/__init__.py b/cagnet/__init__.py new file mode 100644 index 0000000..234bacc --- /dev/null +++ b/cagnet/__init__.py @@ -0,0 +1 @@ +__all__ = [ "nn" ] 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/examples/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py similarity index 100% rename from examples/gcn_conv.py rename to cagnet/nn/conv/gcn_conv.py diff --git a/examples/gcn.py b/examples/gcn.py index d6edda7..3e56a9a 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -14,7 +14,9 @@ import torch_geometric.transforms as T from torch_geometric.utils import add_remaining_self_loops import torch_sparse -from gcn_conv import GCNConv + +from cagnet.nn.conv import GCNConv + import socket class GCN(nn.Module): From e9af1ce40ecb1ee7ece526da441aaec26cbdf4c8 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 10 Jan 2022 21:10:55 -0800 Subject: [PATCH 004/178] Fixed SpMM, added Reddit, and added NCCL --- examples/gcn.py | 44 ++++++++++++++++---------- sparse-extension/sparse_coo_tensor.cpp | 14 ++++---- 2 files changed, 34 insertions(+), 24 deletions(-) diff --git a/examples/gcn.py b/examples/gcn.py index 3e56a9a..178566b 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -10,7 +10,7 @@ import torch.nn.functional as F import torch.autograd.profiler as profiler import torch_geometric -from torch_geometric.datasets import Planetoid, PPI +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 @@ -198,20 +198,24 @@ def evaluate(model, graph, features, labels, nid, nid_count, ampbyp, group): def main(args): # load and preprocess dataset device = torch.device('cuda') - graphname = "Cora" - 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 - adj_matrix = edge_index + + 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(): @@ -270,6 +274,8 @@ def main(args): 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 @@ -277,7 +283,7 @@ def main(args): loss = F.cross_entropy(logits[rank_train_nids], labels_rank[rank_train_nids], reduction="sum") loss_recv = [] - for i in range(size // args.replication): + for i in range(size): loss_recv.append(torch.cuda.FloatTensor(loss.size())) dist.all_gather(loss_recv, loss, group) loss_recv[rank] = loss @@ -297,6 +303,8 @@ def main(args): "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)) @@ -325,6 +333,8 @@ def main(args): 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, @@ -347,7 +357,7 @@ def main(args): 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='gloo', type=str, + 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') diff --git a/sparse-extension/sparse_coo_tensor.cpp b/sparse-extension/sparse_coo_tensor.cpp index 5edf4a8..e6448b1 100644 --- a/sparse-extension/sparse_coo_tensor.cpp +++ b/sparse-extension/sparse_coo_tensor.cpp @@ -155,7 +155,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, cusparseSpMatDescr_t matA; CHECK_CUSPARSE(cusparseCreateCsr(&matA, n, // rows - b_col, // cols + m, // cols nnz, // nnz d_a_csrrows, // csrRowOffsets A_colindices.data(), // csrColInd @@ -167,9 +167,9 @@ void spmm_gpu(const at::Tensor& A_rowindices, cusparseDnMatDescr_t matB; CHECK_CUSPARSE(cusparseCreateDnMat(&matB, - B.size(1), // rows - b_col, // cols - B.size(1), // ld + b_col, // rows + b_row, // cols + b_col, // ld B.data(), // values CUDA_R_32F, // valueType CUSPARSE_ORDER_COL)); // order @@ -182,7 +182,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, cusparseDnMatDescr_t matC; CHECK_CUSPARSE(cusparseCreateDnMat(&matC, n, // rows - B.size(1), // cols + b_col, // cols n, // ld C.data(), // values CUDA_R_32F, // valueType @@ -191,7 +191,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, size_t bufferSize; CHECK_CUSPARSE(cusparseSpMM_bufferSize(handle, // handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA - CUSPARSE_OPERATION_TRANSPOSE, // opB + CUSPARSE_OPERATION_TRANSPOSE, // opB &alpha, // alpha matA, // matA matB, // matB @@ -207,7 +207,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, CHECK_CUSPARSE(cusparseSpMM(handle, // handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA - CUSPARSE_OPERATION_TRANSPOSE, // opB + CUSPARSE_OPERATION_TRANSPOSE, // opB &alpha, // alpha matA, // matA matB, // matB From 1d23259d6f20a9bf3c39134de593a877af2abf80 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 11 Jan 2022 12:07:51 -0800 Subject: [PATCH 005/178] Added partitioning enum --- cagnet/__init__.py | 2 +- cagnet/nn/conv/gcn_conv.py | 9 +++++++-- cagnet/partitionings/__init__.py | 7 +++++++ cagnet/partitionings/partitionings.py | 7 +++++++ examples/gcn.py | 11 +++++++---- 5 files changed, 29 insertions(+), 7 deletions(-) create mode 100644 cagnet/partitionings/__init__.py create mode 100644 cagnet/partitionings/partitionings.py diff --git a/cagnet/__init__.py b/cagnet/__init__.py index 234bacc..efdcbe4 100644 --- a/cagnet/__init__.py +++ b/cagnet/__init__.py @@ -1 +1 @@ -__all__ = [ "nn" ] +__all__ = [ "nn", "partitionings" ] diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index 6d70180..170b6f1 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -3,6 +3,7 @@ import torch.distributed as dist import torch.nn as nn +from cagnet.partitionings import Partitioning from sparse_coo_tensor_cpp import spmm_gpu def broad_func(self, graph, ampbyp, inputs): @@ -35,16 +36,20 @@ def outer_product(mata, matb, group): return matc class GCNConv(nn.Module): - def __init__(self, in_feats, out_feats, device): + def __init__(self, in_feats, out_feats, partitioning, device): 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 def forward(self, gcn, graph, ampbyp, inputs): - return GCNFunc.apply(gcn, graph, ampbyp, inputs, self.weight) + if self.partitioning == Partitioning.ONED: + return GCNFunc.apply(gcn, graph, ampbyp, inputs, self.weight) + else: + print(f"self.partitioning: {self.partitioning}") class GCNFunc(torch.autograd.Function): 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..ec2bebd --- /dev/null +++ b/cagnet/partitionings/partitionings.py @@ -0,0 +1,7 @@ +import enum + +class Partitioning(enum.Enum): + ONED = 0 + ONE5D = 1 + TWOD = 2 + THREED = 3 diff --git a/examples/gcn.py b/examples/gcn.py index 178566b..2725802 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -16,26 +16,28 @@ import torch_sparse from cagnet.nn.conv import GCNConv +from cagnet.partitionings import Partitioning import socket class GCN(nn.Module): - def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, group, device): + def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, group, partitioning, device): 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.device)) + 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.device)) + self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) # output layer - self.layers.append(GCNConv(n_hidden, n_classes, self.device)) + self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) # # input layer # self.mlp_layers.append(GCNConvMLP(in_feats, n_hidden, self.device)) @@ -254,6 +256,7 @@ def main(args): rank, size, group, + Partitioning.ONED, device) # use optimizer From 1e34ac12383899b6ed3d1cc8629122aaf8306b08 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 11 Jan 2022 18:34:06 -0800 Subject: [PATCH 006/178] Added partitioning enum, functional file, and 1.5D refactored --- cagnet/nn/conv/gcn_conv.py | 87 +++++- cagnet/nn/functional.py | 24 ++ cagnet/partitionings/partitionings.py | 2 +- examples/gcn.py | 37 +-- examples/gcn_15d.py | 401 ++++++++++++++++++++++++++ 5 files changed, 521 insertions(+), 30 deletions(-) create mode 100644 cagnet/nn/functional.py create mode 100644 examples/gcn_15d.py diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index 170b6f1..c8f128b 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -6,7 +6,7 @@ from cagnet.partitionings import Partitioning from sparse_coo_tensor_cpp import spmm_gpu -def broad_func(self, graph, ampbyp, inputs): +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) @@ -29,6 +29,47 @@ def broad_func(self, graph, ampbyp, inputs): 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 outer_product(mata, matb, group): matc = torch.mm(mata, matb) dist.all_reduce(matc, op=dist.reduce_op.SUM, group=group) @@ -47,19 +88,53 @@ def __init__(self, in_feats, out_feats, partitioning, device): def forward(self, gcn, graph, ampbyp, inputs): if self.partitioning == Partitioning.ONED: - return GCNFunc.apply(gcn, graph, ampbyp, inputs, self.weight) + return GCNFuncONED.apply(gcn, graph, ampbyp, inputs, self.weight) + elif self.partitioning == Partitioning.ONE5D: + return GCNFuncONE5D.apply(gcn, graph, ampbyp, inputs, self.weight) else: - print(f"self.partitioning: {self.partitioning}") + print(f"self.partitioning: {self.partitioning} one5d: {Partitioning.ONE5D}") -class GCNFunc(torch.autograd.Function): +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(self, graph, ampbyp, inputs) + z = broad_func_one5d(self, graph, ampbyp, inputs) z = z.mm(weight) ctx.save_for_backward(inputs, weight) @@ -77,7 +152,7 @@ def backward(ctx, grad_output): self = ctx.self # Assumes graph is undirected and A = A^T - ag = broad_func(self, graph, ampbyp, grad_output) + ag = broad_func_one5d(self, graph, ampbyp, grad_output) grad_input = ag.mm(weight.t()) grad_weight = outer_product(inputs.t(), ag, self.group) diff --git a/cagnet/nn/functional.py b/cagnet/nn/functional.py new file mode 100644 index 0000000..e56a4b4 --- /dev/null +++ b/cagnet/nn/functional.py @@ -0,0 +1,24 @@ +from cagnet.partitionings import Partitioning +import torch +import torch.nn.functional as F +import torch.distributed as dist + +def relu(x, partitioning): + if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: + return F.relu(x) + +def log_softmax(x, partitioning, dim=1): + if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: + return F.log_softmax(x, dim) + +def cross_entropy(logits_rank, labels_rank, node_count, partitioning, rank, group, size): + 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 diff --git a/cagnet/partitionings/partitionings.py b/cagnet/partitionings/partitionings.py index ec2bebd..1dbd3ab 100644 --- a/cagnet/partitionings/partitionings.py +++ b/cagnet/partitionings/partitionings.py @@ -1,6 +1,6 @@ import enum -class Partitioning(enum.Enum): +class Partitioning(enum.IntEnum): ONED = 0 ONE5D = 1 TWOD = 2 diff --git a/examples/gcn.py b/examples/gcn.py index 2725802..879f461 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -7,7 +7,6 @@ import torch import torch.distributed as dist import torch.nn as nn -import torch.nn.functional as F import torch.autograd.profiler as profiler import torch_geometric from torch_geometric.datasets import Planetoid, Reddit @@ -17,11 +16,13 @@ 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, group, partitioning, device): + 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 @@ -39,22 +40,14 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, group, p # output layer self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) - # # input layer - # self.mlp_layers.append(GCNConvMLP(in_feats, n_hidden, self.device)) - # # hidden layers - # for i in range(n_layers - 1): - # self.mlp_layers.append(GCNConvMLP(n_hidden, n_hidden, self.device)) - # # output layer - # self.mlp_layers.append(GCNConvMLP(n_hidden, n_classes, self.device)) - def forward(self, graph, inputs, ampbyp): h = inputs for l, layer in enumerate(self.layers): h = layer(self, graph, ampbyp, h) if l != len(self.layers) - 1: - h = F.relu(h) + h = CAGF.relu(h, self.partitioning) - h = F.log_softmax(h, dim=1) + h = CAGF.log_softmax(h, self.partitioning, dim=1) return h # Normalize all elements according to KW's normalization rule @@ -237,6 +230,8 @@ def main(args): if args.normalize: adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + partitioning = Partitioning.ONED + 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) @@ -255,9 +250,9 @@ def main(args): args.n_layers, rank, size, - group, - Partitioning.ONED, - device) + partitioning, + device, + group) # use optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) @@ -283,14 +278,8 @@ def main(args): # forward logits = model(g_loc, features_loc, ampbyp) - loss = F.cross_entropy(logits[rank_train_nids], labels_rank[rank_train_nids], 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) / train_nid.size(0) + loss = CAGF.cross_entropy(logits[rank_train_nids], labels_rank[rank_train_nids], train_nid.size(0), \ + partitioning, rank, group, size) optimizer.zero_grad() loss.backward() @@ -366,6 +355,8 @@ def main(args): 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) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py new file mode 100644 index 0000000..431a1b8 --- /dev/null +++ b/examples/gcn_15d.py @@ -0,0 +1,401 @@ +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.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, replication, 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.row_groups = row_groups + self.col_groups = col_groups + self.device = device + self.partitioning = partitioning + self.replication = replication + 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, ampbyp, h) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + + h = CAGF.log_softmax(h, self.partitioning, dim=1) + return h + +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].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 one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, normalize, 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], 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) + + 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 / 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 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.ONE5D + + row_groups, col_groups = get_proc_groups(rank, size, args.replication) + + rank_c = rank // args.replication + rank_col = rank % args.replication + if rank_c >= (size // args.replication): + return + + features_loc, g_loc, ampbyp = one5d_partition(rank, size, inputs, adj_matrix, data, \ + inputs, num_classes, args.replication, args.normalize, \ + args.normalize) + + 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, + args.replication, + device, + row_groups=row_groups, + col_groups=col_groups) + + # use optimizer + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) + + rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] + rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] + labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] + rank_train_nids = rank_train_mask.nonzero().squeeze() + rank_test_nids = rank_test_mask.nonzero().squeeze() + + 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, ampbyp) + loss = CAGF.cross_entropy(logits[rank_train_nids], labels_rank[rank_train_nids], train_nid.size(0), \ + partitioning, rank_c, col_groups[0], size // args.replication) + + 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='ONE5D', type=str, + help='partitioning strategy to use') + parser.add_argument('--replication', default=1, type=int, + help='partitioning strategy to use') + args = parser.parse_args() + print(args) + + main(args) From e3ef8d2a7d6fd0fe507f6985b18030311c61867c Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 11 Jan 2022 21:49:54 -0800 Subject: [PATCH 007/178] Added refactored 2D implementation --- cagnet/nn/conv/gcn_conv.py | 325 ++++++++++++++++++++++++++- cagnet/nn/functional.py | 179 ++++++++++++++- examples/gcn.py | 6 +- examples/gcn_15d.py | 7 +- examples/gcn_2d.py | 446 +++++++++++++++++++++++++++++++++++++ 5 files changed, 951 insertions(+), 12 deletions(-) create mode 100644 examples/gcn_2d.py diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index c8f128b..0bb1304 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -4,7 +4,7 @@ import torch.nn as nn from cagnet.partitionings import Partitioning -from sparse_coo_tensor_cpp import spmm_gpu +from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu def broad_func_oned(self, graph, ampbyp, inputs): n_per_proc = math.ceil(float(graph.size(0) / self.size)) @@ -70,6 +70,200 @@ def broad_func_one5d(self, graph, ampbyp, inputs): 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): + 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 outer_product(mata, matb, group): matc = torch.mm(mata, matb) dist.all_reduce(matc, op=dist.reduce_op.SUM, group=group) @@ -86,11 +280,13 @@ def __init__(self, in_feats, out_feats, partitioning, device): self.weight = nn.Parameter(weight_nonleaf) self.partitioning = partitioning - def forward(self, gcn, graph, ampbyp, inputs): + def forward(self, gcn, graph, inputs, 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) else: print(f"self.partitioning: {self.partitioning} one5d: {Partitioning.ONE5D}") @@ -158,3 +354,128 @@ def backward(ctx, grad_output): 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 diff --git a/cagnet/nn/functional.py b/cagnet/nn/functional.py index e56a4b4..5298435 100644 --- a/cagnet/nn/functional.py +++ b/cagnet/nn/functional.py @@ -1,17 +1,27 @@ 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): + return math.floor(math.sqrt(size)) + +def proc_col_size(size): + return math.floor(math.sqrt(size)) + def relu(x, partitioning): - if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: - return F.relu(x) + return F.relu(x) -def log_softmax(x, partitioning, dim=1): +def log_softmax(self, x, partitioning, dim=1): if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: 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): -def cross_entropy(logits_rank, labels_rank, node_count, partitioning, rank, group, size): if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: loss = F.cross_entropy(logits_rank, labels_rank, reduction="sum") @@ -22,3 +32,164 @@ def cross_entropy(logits_rank, labels_rank, node_count, partitioning, rank, grou loss_recv[rank] = loss loss = sum(loss_recv) / node_count return loss + + elif partitioning == Partitioning.TWOD: + proc_row = proc_row_size(size) + proc_col = proc_col_size(size) + + 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 + + +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/examples/gcn.py b/examples/gcn.py index 879f461..e5c244b 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -43,11 +43,11 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, partitio def forward(self, graph, inputs, ampbyp): h = inputs for l, layer in enumerate(self.layers): - h = layer(self, graph, ampbyp, h) + h = layer(self, graph, h, ampbyp) if l != len(self.layers) - 1: h = CAGF.relu(h, self.partitioning) - h = CAGF.log_softmax(h, self.partitioning, dim=1) + h = CAGF.log_softmax(self, h, self.partitioning, dim=1) return h # Normalize all elements according to KW's normalization rule @@ -279,7 +279,7 @@ def main(args): # 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), \ - partitioning, rank, group, size) + num_classes, partitioning, rank, group, size) optimizer.zero_grad() loss.backward() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 431a1b8..4ecd026 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -46,11 +46,11 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, partitio def forward(self, graph, inputs, ampbyp): h = inputs for l, layer in enumerate(self.layers): - h = layer(self, graph, ampbyp, h) + h = layer(self, graph, h, ampbyp) if l != len(self.layers) - 1: h = CAGF.relu(h, self.partitioning) - h = CAGF.log_softmax(h, self.partitioning, dim=1) + h = CAGF.log_softmax(self, h, self.partitioning, dim=1) return h def get_proc_groups(rank, size, replication): @@ -315,7 +315,8 @@ def main(args): # 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), \ - partitioning, rank_c, col_groups[0], size // args.replication) + num_classs, partitioning, rank_c, col_groups[0], \ + size // args.replication) optimizer.zero_grad() loss.backward() diff --git a/examples/gcn_2d.py b/examples/gcn_2d.py new file mode 100644 index 0000000..783492d --- /dev/null +++ b/examples/gcn_2d.py @@ -0,0 +1,446 @@ +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) + proc_col = CAGF.proc_col_size(size) + 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) + proc_col = CAGF.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): + 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)))) + + proc_row = CAGF.proc_row_size(size) + proc_col = CAGF.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] + + 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) + proc_col = CAGF.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 + +def twod_partition(rank, size, inputs, adj_matrix, data, features, classes, normalize, device): + node_count = inputs.size(0) + proc_row = CAGF.proc_row_size(size) + proc_col = CAGF.proc_col_size(size) + + 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) + proc_col = CAGF.proc_col_size(size) + rank_row = int(rank / proc_col) + rank_col = rank % proc_col + + features_loc, g_loc = twod_partition(rank, size, inputs, adj_matrix, data, \ + inputs, num_classes, args.normalize, args.normalize) + + 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) From 387d9f78654a578793985068d6f1ed11c3f4903b Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 13 Jan 2022 12:30:25 -0800 Subject: [PATCH 008/178] Refactored 3D implementation --- .gitignore | 2 + cagnet/nn/conv/gcn_conv.py | 466 +++++++++++++++++++++++++++++++- cagnet/nn/functional.py | 89 ++++++- examples/gcn.py | 55 ++-- examples/gcn_15d.py | 2 +- examples/gcn_2d.py | 29 +- examples/gcn_3d.py | 533 +++++++++++++++++++++++++++++++++++++ gcn_distr.py | 7 +- gcn_distr_3d.py | 17 +- 9 files changed, 1146 insertions(+), 54 deletions(-) create mode 100644 examples/gcn_3d.py diff --git a/.gitignore b/.gitignore index 5ed031f..0081ad0 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,5 @@ sparse_coo_tensor_cpp.egg-info/ .nfs* *.txt *.pt +slurm_outputs/ +job-* diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index 0bb1304..f257592 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -235,7 +235,13 @@ def summa(self, adj_matrix, inputs, height, middim, width): return z_loc -def transpose(self, mat, height, width): +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 @@ -264,12 +270,286 @@ def transpose(self, mat, height, width): 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): super(GCNConv, self).__init__() @@ -287,8 +567,8 @@ def forward(self, gcn, graph, inputs, ampbyp=None): return GCNFuncONE5D.apply(gcn, graph, ampbyp, inputs, self.weight) elif self.partitioning == Partitioning.TWOD: return GCNFuncTWOD.apply(gcn, graph, inputs, self.weight) - else: - print(f"self.partitioning: {self.partitioning} one5d: {Partitioning.ONE5D}") + elif self.partitioning == Partitioning.THREED: + return GCNFuncTHREED.apply(gcn, graph, inputs, self.weight) class GCNFuncONED(torch.autograd.Function): @@ -296,7 +576,7 @@ class GCNFuncONED(torch.autograd.Function): def forward(ctx, self, graph, ampbyp, inputs, weight): # inputs: H # graph: A - # weight; W + # weight: W z = broad_func_oned(self, graph, ampbyp, inputs) z = z.mm(weight) @@ -328,7 +608,7 @@ class GCNFuncONE5D(torch.autograd.Function): def forward(ctx, self, graph, ampbyp, inputs, weight): # inputs: H # graph: A - # weight; W + # weight: W z = broad_func_one5d(self, graph, ampbyp, inputs) z = z.mm(weight) @@ -360,7 +640,7 @@ class GCNFuncTWOD(torch.autograd.Function): def forward(ctx, self, graph, inputs, weight): # inputs: H # graph: A - # weight; W + # weight: W graph_t = graph # Only true for undirected graphs @@ -438,7 +718,7 @@ def backward(ctx, grad_output): # Collect grad_weight's across processes grad_weight_recv = [] - max_row_chunk = max(chunk_sizes_col) #transpose + 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( @@ -479,3 +759,175 @@ def backward(ctx, grad_output): 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 index 5298435..0d9a139 100644 --- a/cagnet/nn/functional.py +++ b/cagnet/nn/functional.py @@ -4,11 +4,39 @@ import torch.nn.functional as F import torch.distributed as dist -def proc_row_size(size): - return math.floor(math.sqrt(size)) - -def proc_col_size(size): - return math.floor(math.sqrt(size)) +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) @@ -34,8 +62,8 @@ def cross_entropy(logits_rank, labels_rank, node_count, total_classes, partition return loss elif partitioning == Partitioning.TWOD: - proc_row = proc_row_size(size) - proc_col = proc_col_size(size) + 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 @@ -79,6 +107,53 @@ def cross_entropy(logits_rank, labels_rank, node_count, total_classes, partition 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 diff --git a/examples/gcn.py b/examples/gcn.py index e5c244b..f12919c 100644 --- a/examples/gcn.py +++ b/examples/gcn.py @@ -9,6 +9,7 @@ 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 @@ -196,21 +197,37 @@ def main(args): 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 + 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(): @@ -232,9 +249,11 @@ def main(args): 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) @@ -260,13 +279,10 @@ def main(args): n_per_proc = math.ceil(float(inputs.size(0)) / size) rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank] - rank_test_mask = torch.split(data.test_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() - rank_test_nids = rank_test_mask.nonzero().squeeze() train_nid = data.train_mask.nonzero().squeeze() - test_nid = data.test_mask.nonzero().squeeze() torch.manual_seed(0) total_start = time.time() @@ -295,8 +311,9 @@ def main(args): "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"total_time: {total_stop - total_start}") + print(f"rank: {rank} total_time: {total_stop - total_start}") """ # print(prof.key_averages().table(sort_by='self_cpu_time_total', row_limit=10)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 4ecd026..791e064 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -315,7 +315,7 @@ def main(args): # 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_classs, partitioning, rank_c, col_groups[0], \ + num_classes, partitioning, rank_c, col_groups[0], \ size // args.replication) optimizer.zero_grad() diff --git a/examples/gcn_2d.py b/examples/gcn_2d.py index 783492d..b8416e2 100644 --- a/examples/gcn_2d.py +++ b/examples/gcn_2d.py @@ -35,8 +35,8 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, node_count, rank, si self.n_classes = n_classes self.timings = dict() - proc_row = CAGF.proc_row_size(size) - proc_col = CAGF.proc_col_size(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 @@ -64,8 +64,8 @@ def forward(self, graph, inputs, ampbyp=None): return h def get_proc_groups(rank, size): - proc_row = CAGF.proc_row_size(size) - proc_col = CAGF.proc_col_size(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 @@ -80,10 +80,6 @@ def get_proc_groups(rank, size): for i in range(proc_col): col_groups.append(dist.new_group(list(range(i, size, proc_row)))) - proc_row = CAGF.proc_row_size(size) - proc_col = CAGF.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: @@ -150,8 +146,8 @@ def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx, normaliza # 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) - proc_col = CAGF.proc_col_size(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] @@ -166,10 +162,10 @@ def split_coo(adj_matrix, node_count, n_per_proc, dim, size): return am_partitions, vtx_indices -def twod_partition(rank, size, inputs, adj_matrix, data, features, classes, normalize, device): +def twod_partition(rank, size, inputs, adj_matrix, normalize, device): node_count = inputs.size(0) - proc_row = CAGF.proc_row_size(size) - proc_col = CAGF.proc_col_size(size) + 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")) @@ -313,13 +309,12 @@ def main(args): row_groups, col_groups, transpose_group = get_proc_groups(rank, size) - proc_row = CAGF.proc_row_size(size) - proc_col = CAGF.proc_col_size(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, data, \ - inputs, num_classes, args.normalize, args.normalize) + 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) 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/gcn_distr.py b/gcn_distr.py index d982788..bf88335 100644 --- a/gcn_distr.py +++ b/gcn_distr.py @@ -676,7 +676,7 @@ def main(): global device global graphname - print(socket.gethostname()) + print(socket.gethostname(), flush=True) seed = 0 if not download: @@ -699,7 +699,7 @@ def main(): dist.init_process_group(backend='nccl') rank = dist.get_rank() size = dist.get_world_size() - print("Processes: " + str(size)) + print("Processes: " + str(size), flush=True) # device = torch.device('cpu') devid = rank_to_devid(rank, acc_per_rank) @@ -734,6 +734,8 @@ def main(): edge_index = data.edge_index num_features = dataset.num_features num_classes = dataset.num_classes + print(f"edge_index.size: {edge_index.size()}") + print(f"edge_index: {edge_index}") elif graphname == 'Amazon': # path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname) # edge_index = torch.load(path + "/processed/amazon_graph.pt") @@ -752,6 +754,7 @@ def main(): data.y = torch.rand(n).uniform_(0, num_classes - 1).long() data.train_mask = torch.ones(n).long() # edge_index = edge_index.to(device) + edge_index = edge_index.t_() print(f"edge_index.size: {edge_index.size()}", flush=True) print(f"edge_index: {edge_index}", flush=True) data = data.to(device) diff --git a/gcn_distr_3d.py b/gcn_distr_3d.py index 20bd16b..d437d51 100644 --- a/gcn_distr_3d.py +++ b/gcn_distr_3d.py @@ -1270,6 +1270,9 @@ def proc_c_size(size): print(f"CUBE ROOT ERROR") def twod_partition(rank, size, inputs, adj_matrix, data, features, classes, device): + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) + node_count = inputs.size(0) proc_row = proc_row_size(size) proc_col = proc_col_size(size) @@ -1664,7 +1667,7 @@ def main(): # mid_layer = 16 if graphname == 'Cora': - dataset = Planetoid(path, graphname, T.NormalizeFeatures()) + dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures()) data = dataset[0] num_features = dataset.num_features num_classes = dataset.num_classes @@ -1719,6 +1722,18 @@ def main(): 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() From 767159b364b59b66bff56c315753b55fa0ee854f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 13 Jan 2022 12:41:36 -0800 Subject: [PATCH 009/178] Removed old gcn_distr files --- .gitignore | 4 + gcn_distr.py | 837 --------------------- gcn_distr_15d.py | 870 ---------------------- gcn_distr_2d.py | 1575 --------------------------------------- gcn_distr_3d.py | 1822 ---------------------------------------------- 5 files changed, 4 insertions(+), 5104 deletions(-) delete mode 100644 gcn_distr.py delete mode 100644 gcn_distr_15d.py delete mode 100644 gcn_distr_2d.py delete mode 100644 gcn_distr_3d.py diff --git a/.gitignore b/.gitignore index 0081ad0..12139a6 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,7 @@ sparse_coo_tensor_cpp.egg-info/ *.pt slurm_outputs/ job-* +gcn_distr.py +gcn_distr_15d.py +gcn_distr_2d.py +gcn_distr_3d.py diff --git a/gcn_distr.py b/gcn_distr.py deleted file mode 100644 index bf88335..0000000 --- a/gcn_distr.py +++ /dev/null @@ -1,837 +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(), flush=True) - 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), flush=True) - - # 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 - print(f"edge_index.size: {edge_index.size()}") - print(f"edge_index: {edge_index}") - 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) - edge_index = edge_index.t_() - 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 d437d51..0000000 --- a/gcn_distr_3d.py +++ /dev/null @@ -1,1822 +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): - inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) - - 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, 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 - 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"] - - # 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)) - # 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()) From 8d698a738d20358f39048a89950440dfce608904 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 13 Jan 2022 12:43:08 -0800 Subject: [PATCH 010/178] Removed old gcn.py and reddit.py files --- .gitignore | 2 + examples/{gcn.py => gcn_1d.py} | 0 gcn.py | 126 --------------------------------- reddit.py | 70 ------------------ 4 files changed, 2 insertions(+), 196 deletions(-) rename examples/{gcn.py => gcn_1d.py} (100%) delete mode 100644 gcn.py delete mode 100644 reddit.py diff --git a/.gitignore b/.gitignore index 12139a6..f776463 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,5 @@ gcn_distr.py gcn_distr_15d.py gcn_distr_2d.py gcn_distr_3d.py +gcn.py +reddit.py diff --git a/examples/gcn.py b/examples/gcn_1d.py similarity index 100% rename from examples/gcn.py rename to examples/gcn_1d.py 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/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__) From d00b85288f4cfb7d513bf2d4578d491ce79a21cf Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 13 Jan 2022 13:27:40 -0800 Subject: [PATCH 011/178] Added requirements.txt --- .gitignore | 1 + README.md | 2 +- requirements.txt | 151 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 153 insertions(+), 1 deletion(-) create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore index f776463..a7ecc92 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ sparse_coo_tensor_cpp.egg-info/ .*.swp .nfs* *.txt +!requirements.txt *.pt slurm_outputs/ job-* diff --git a/README.md b/README.md index db6581f..675bcd1 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ For more information, please read our ACM/IEEE SC'20 paper [Reducing Communicati **Contact:** Alok Tripathy () ## Dependencies -- Python 3.6.10 +- Python 3.8.11 - PyTorch 1.3.1 - PyTorch Geometric (PyG) 1.3.2 - CUDA 10.1 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 From ab5312c733e3e3d79a2d5c11c43fd9fe51afc8d6 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 13 Jan 2022 13:43:43 -0800 Subject: [PATCH 012/178] Updated README for Perlmutter --- README.md | 52 ++++++++++++++++++++++++++-------------------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index 675bcd1..20821c5 100644 --- a/README.md +++ b/README.md @@ -11,10 +11,16 @@ For more information, please read our ACM/IEEE SC'20 paper [Reducing Communicati ## Dependencies - Python 3.8.11 -- PyTorch 1.3.1 -- PyTorch Geometric (PyG) 1.3.2 -- CUDA 10.1 -- GCC 6.4.0 +- PyTorch 1.9.0 +- PyTorch Geometric (PyG) 1.7.2 +- CUDA 10.2 +- GCC 9.3.0 + +On NERSC Perlmutter, all of these dependencies can be accessed with the following +```bash +module load tudatoolkit/21.3_10.2 # CUDA 10.2 +module load pytorch/1.9.0 # PyTorch and PyG +``` On OLCF Summit, all of these dependencies can be accessed with the following ```bash @@ -43,45 +49,39 @@ python setup.py install ## Documentation Each algorithm in CAGNET is implemented in a separate file. -- `gcn_distr.py` : 1D algorithm -- `gcn_distr_15d.py` : 1.5D algorithm -- `gcn_distr_2d.py` : 2D algorithm -- `gcn_distr_3d.py` : 3D algorithm +- `examples/gcn_1d.py` : 1D algorithm +- `examples/gcn_15d.py` : 1.5D algorithm +- `examples/gcn_2d.py` : 2D algorithm +- `examples/gcn_3d.py` : 3D algorithm Each file also as the following flags: -- `--accperrank ` : Number of GPUs on each node -- `--epochs ` : Number of epochs to run training -- `--graphname ` : Graph dataset to run training on +- `--n-epochs ` : Number of epochs to run training +- `--dataset ` : Graph dataset to run training on - `--timing ` : Enable timing barriers to time phases in training -- `--midlayer ` : Number of activations in the hidden layer -- `--runcount ` : Number of times to run training -- `--normalization ` : Normalize adjacency matrix in preprocessing -- `--activations ` : Enable activation functions between layers +- `--n-hidden ` : Number of activations in the hidden layer +- `--normalize` : Normalize adjacency matrix in preprocessing - `--accuracy ` : Compute and print accuracy metrics (Reddit only) - `--replication ` : Replication factor (1.5D algorithm only) -- `--download ` : Download the Reddit dataset Some of these flags do not currently exist for the 3D algorithm. -Amazon/Protein datasets must exist as COO files in `../data//processed/`, compressed with pickle. -For Reddit, PyG handles downloading and accessing the dataset (see below). +Amazon/Protein datasets must exist as COO files in `../../data//processed/`, compressed with pickle. +For Cora and Reddit, PyG handles downloading and accessing the dataset. -## Running on OLCF Summit (example) +## Running on NERSC Perlmutter (example) To run the CAGNET 1.5D algorithm on Reddit with -- 16 processes +- 4 processes - 100 epochs - 16 hidden layer activations - 2-factor replication +- normalization +- no weight decay -run the following command to download the Reddit dataset: - -`python gcn_distr_15d.py --graphname=Reddit --download=True` - -This will download Reddit into `../data`. After downloading the Reddit dataset, run the following command to run training +run the following command with Slurm: -`ddlrun -x WORLD_SIZE=16 -x MASTER_ADDR=$(echo $LSB_MCPU_HOSTS | cut -d " " -f 3) -x MASTER_PORT=1234 -accelerators 6 python gcn_distr_15d.py --accperrank=6 --epochs=100 --graphname=Reddit --timing=False --midlayer=16 --runcount=1 --replication=2` +`cd examples && srun -n 4 --ntasks-per-node=4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset=Reddit --n-epochs=100 --n-hidden 16 --weight-decay 0 --replication 2 --normalize` ## Citation From badf6cc71d1d05c9672b4d60df0582f83f758015 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 13 Jan 2022 13:44:36 -0800 Subject: [PATCH 013/178] Updated README --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 20821c5..33f8685 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ module load tudatoolkit/21.3_10.2 # CUDA 10.2 module load pytorch/1.9.0 # PyTorch and PyG ``` -On OLCF Summit, all of these dependencies can be accessed with the following +On OLCF Summit, older versions of these dependencies can be accessed with the following ```bash module load cuda # CUDA 10.1 module load gcc # GCC 6.4.0 From dfc48e16634e033224f197ec9a5c5e3504de48fa Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 19 Jan 2022 18:54:12 -0800 Subject: [PATCH 014/178] Initial LADIES sampler implementation --- .gitignore | 1 + cagnet/samplers/__init__.py | 7 +++ cagnet/samplers/ladies.py | 67 +++++++++++++++++++++++ examples/gcn_15d.py | 106 ++++++++++++++++++++++++++---------- 4 files changed, 151 insertions(+), 30 deletions(-) create mode 100644 cagnet/samplers/__init__.py create mode 100644 cagnet/samplers/ladies.py diff --git a/.gitignore b/.gitignore index a7ecc92..3c5cbdf 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,4 @@ gcn_distr_2d.py gcn_distr_3d.py gcn.py reddit.py +tests/ diff --git a/cagnet/samplers/__init__.py b/cagnet/samplers/__init__.py new file mode 100644 index 0000000..fd68782 --- /dev/null +++ b/cagnet/samplers/__init__.py @@ -0,0 +1,7 @@ +from .ladies import ladies_sampler + +__all__ = [ + 'ladies_sampler', +] + +classes = __all__ diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py new file mode 100644 index 0000000..3da7fa4 --- /dev/null +++ b/cagnet/samplers/ladies.py @@ -0,0 +1,67 @@ +import torch +import torch_sparse + +def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, train_nodes): + frontiers = torch.cuda.IntTensor(n_layers, batch_size, mb_count) + node_count = adj_matrix.size(0) + + torch.manual_seed(0) + for i in range(mb_count): + idx = torch.randperm(train_nodes.size(0))[:batch_size] + frontiers[0,:,i] = train_nodes[idx] + + for i in range(1, n_layers + 1): + nnz = frontiers[i - 1, :, 0].size(0) + + # A * Q^i + # indices would change based on mb_count + frontier_indices = torch.stack((torch.cuda.IntTensor(batch_size).fill_(0), frontiers[i - 1,:,0])) + frontier_values = torch.cuda.FloatTensor(nnz).fill_(1.0) + adj_mat_squared = torch.pow(adj_matrix._values(), 2) + p_num_indices, p_num_values = torch_sparse.spspmm(frontier_indices.long(), frontier_values, + adj_matrix._indices().long(), adj_mat_squared, + mb_count, node_count, node_count, coalesced=True) + + p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count)) + p_den = torch.sparse.sum(p_num, dim=1) + + # Assumes mb_count = 1 + p_den = p_den._values().item() + p = torch.div(p_num, p_den) + + print(f"p: {p}") + print(f"p.sum: {torch.sparse.sum(p)}") + + print(f"neighbors: {p._indices()[1,:]}") + + max_prob = torch.max(p._values()) + frontier_next = torch.cuda.LongTensor(frontier_size).fill_(node_count) + sampled_count = 0 + + while sampled_count < frontier_size: + # sample u.a.r. sampled_count vertices per minibatch from vertices in the minibatch's neighborhood + # neighbored is vertices with nonzero probabilities + # TODO: What if a vertex is sampled more than once (idk if w/o replacement is possible) + sampled_verts_ids = torch.randint(p._nnz(), (frontier_size,)) + + p_loc = p._values()[sampled_verts_ids] + + dart_throws = torch.cuda.FloatTensor(frontier_size).uniform_(to=max_prob) + + dart_misses = p_loc <= dart_throws + dart_hits = p_loc > dart_throws + hit_count = dart_hits.sum() + hit_verts_ids = sampled_verts_ids[dart_hits] + sampled_verts = p._indices()[1][hit_verts_ids] + print(f"sampled_verts: {sampled_verts}") + + frontier_next[dart_hits] = sampled_verts + p._values()[hit_verts_ids] = 0.0 + + sampled_count += hit_count + + print(f"sampled_count: {sampled_count}") + print(f"hit_count: {hit_count}") + + print(f"frontier_next: {frontier_next}") + return p diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 791e064..57c759d 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -16,6 +16,7 @@ from cagnet.nn.conv import GCNConv from cagnet.partitionings import Partitioning +from cagnet.samplers import ladies_sampler import cagnet.nn.functional as CAGF import socket @@ -127,7 +128,28 @@ def split_coo(adj_matrix, node_count, n_per_proc, dim): return am_partitions, vtx_indices -def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, normalize, device): +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 = 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().float(), + mx._indices(), mx._values().float(), + 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, size=(r_mat_inv.size(0), mx.size(1))) + return mx + +def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, \ + normalize): node_count = inputs.size(0) # n_per_proc = math.ceil(float(node_count) / size) n_per_proc = math.ceil(float(node_count) / (size / replication)) @@ -223,21 +245,37 @@ def main(args): 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 + 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(): @@ -254,8 +292,7 @@ def main(args): 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)) + adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) partitioning = Partitioning.ONE5D @@ -267,14 +304,30 @@ def main(args): return features_loc, g_loc, ampbyp = one5d_partition(rank, size, inputs, adj_matrix, data, \ - inputs, num_classes, args.replication, args.normalize, \ + inputs, num_classes, args.replication, \ args.normalize) + g_loc = g_loc.coalesce() + g_loc = row_normalize(g_loc) 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) + n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) + + rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] + rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] + labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] + rank_train_nids = rank_train_mask.nonzero().squeeze() + rank_test_nids = rank_test_mask.nonzero().squeeze() + + train_nid = data.train_mask.nonzero().squeeze() + test_nid = data.test_mask.nonzero().squeeze() + + result = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, args.n_layers, train_nid) + + return result # return here while testing sampling code # create GCN model torch.manual_seed(0) @@ -293,17 +346,6 @@ def main(args): # use optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) - n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) - - rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] - rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] - labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] - rank_train_nids = rank_train_mask.nonzero().squeeze() - rank_test_nids = rank_test_mask.nonzero().squeeze() - - 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): @@ -376,6 +418,10 @@ def main(args): 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=int, default=64, + help="number of vertices per layer of layer-wise minibatch") parser.add_argument("--weight-decay", type=float, default=5e-4, help="Weight for L2 loss") parser.add_argument("--aggregator-type", type=str, default="gcn", From 4496c13ddda88b42aaf85ecf91fb10da7f27f122 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 20 Jan 2022 12:30:13 -0800 Subject: [PATCH 015/178] Fixed sampled_count computation at end of sampling loop --- cagnet/samplers/ladies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 3da7fa4..38c3a5d 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -58,7 +58,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr frontier_next[dart_hits] = sampled_verts p._values()[hit_verts_ids] = 0.0 - sampled_count += hit_count + sampled_count = torch.sum(frontier_next != node_count) print(f"sampled_count: {sampled_count}") print(f"hit_count: {hit_count}") From fa8f2c30ad212c80666b1b1a1816dc1ef570c308 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 21 Jan 2022 13:50:55 -0800 Subject: [PATCH 016/178] Added scaffolding for selection multiplication --- cagnet/samplers/ladies.py | 47 ++++++++++++++++++++++++++++++--------- 1 file changed, 37 insertions(+), 10 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 38c3a5d..e920341 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -2,20 +2,31 @@ import torch_sparse def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, train_nodes): - frontiers = torch.cuda.IntTensor(n_layers, batch_size, mb_count) + # frontiers = torch.cuda.IntTensor(n_layers, batch_size, mb_count) + batches = torch.cuda.IntTensor(batch_size, mb_count) # initially the minibatch + current_frontier = torch.cuda.IntTensor(batch_size + frontier_size, mb_count) + # frontiers = torch.cuda.IntTensor(n_layers - 1, batch_size + frontier_size, mb_count) + adj_matrices = [None] * n_layers node_count = adj_matrix.size(0) torch.manual_seed(0) + # Generate minibatch vertices for i in range(mb_count): idx = torch.randperm(train_nodes.size(0))[:batch_size] - frontiers[0,:,i] = train_nodes[idx] + batches[:,i] = train_nodes[idx] - for i in range(1, n_layers + 1): - nnz = frontiers[i - 1, :, 0].size(0) + for i in range(n_layers): + if i == 0: + nnz = batches[:, 0].size(0) # assumes mb_count == 1 + else: + nnz = current_frontier[:, 0].size(0) # assumes mb_count == 1 # A * Q^i # indices would change based on mb_count - frontier_indices = torch.stack((torch.cuda.IntTensor(batch_size).fill_(0), frontiers[i - 1,:,0])) + if i == 0: + frontier_indices = torch.stack((torch.cuda.IntTensor(nnz).fill_(0), batches[:,0])) + else: + frontier_indices = torch.stack((torch.cuda.IntTensor(nnz).fill_(0), current_frontier[:,0])) frontier_values = torch.cuda.FloatTensor(nnz).fill_(1.0) adj_mat_squared = torch.pow(adj_matrix._values(), 2) p_num_indices, p_num_values = torch_sparse.spspmm(frontier_indices.long(), frontier_values, @@ -35,7 +46,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr print(f"neighbors: {p._indices()[1,:]}") max_prob = torch.max(p._values()) - frontier_next = torch.cuda.LongTensor(frontier_size).fill_(node_count) + next_frontier = torch.cuda.LongTensor(frontier_size).fill_(node_count) sampled_count = 0 while sampled_count < frontier_size: @@ -55,13 +66,29 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr sampled_verts = p._indices()[1][hit_verts_ids] print(f"sampled_verts: {sampled_verts}") - frontier_next[dart_hits] = sampled_verts + next_frontier[dart_hits] = sampled_verts p._values()[hit_verts_ids] = 0.0 - sampled_count = torch.sum(frontier_next != node_count) + sampled_count = torch.sum(next_frontier != node_count) print(f"sampled_count: {sampled_count}") print(f"hit_count: {hit_count}") - print(f"frontier_next: {frontier_next}") - return p + next_frontier = torch.cat((next_frontier, batches[:,0]), dim=0) + + if i == 0: + row_select_mtx_indices = torch.stack((torch.range(nnz), batches[:,0])) + else: + row_select_mtx_indices = torch.stack((torch.range(nnz), current_frontier[:,0])) + row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) + + col_select_mtx_indices = torch.stack((next_frontier, torch.range(next_frontier.size(0))) + col_select_mtx_values = torch.cuda.FloatTensor(next_frontier.size(0)).fill_(1.0) + + # multiply row_select matrix with adj_matrix + # multiply adj_matrix with col_select matrix + # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] + + current_frontier[:,0] = frontier_next + + return adj_matrices From cdaa648fdd49ec87a91ba0f9c7a26a5d315728a4 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 27 Jan 2022 16:20:22 -0800 Subject: [PATCH 017/178] Added test for LADIES based on sample --- cagnet/samplers/ladies.py | 32 +++++++++++++++++--------------- examples/gcn_15d.py | 18 ++++++++++++++++-- 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index e920341..901a7aa 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -40,11 +40,6 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr p_den = p_den._values().item() p = torch.div(p_num, p_den) - print(f"p: {p}") - print(f"p.sum: {torch.sparse.sum(p)}") - - print(f"neighbors: {p._indices()[1,:]}") - max_prob = torch.max(p._values()) next_frontier = torch.cuda.LongTensor(frontier_size).fill_(node_count) sampled_count = 0 @@ -64,31 +59,38 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr hit_count = dart_hits.sum() hit_verts_ids = sampled_verts_ids[dart_hits] sampled_verts = p._indices()[1][hit_verts_ids] - print(f"sampled_verts: {sampled_verts}") next_frontier[dart_hits] = sampled_verts p._values()[hit_verts_ids] = 0.0 sampled_count = torch.sum(next_frontier != node_count) - print(f"sampled_count: {sampled_count}") - print(f"hit_count: {hit_count}") - next_frontier = torch.cat((next_frontier, batches[:,0]), dim=0) - + if i == 0: - row_select_mtx_indices = torch.stack((torch.range(nnz), batches[:,0])) + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches[:,0])) else: - row_select_mtx_indices = torch.stack((torch.range(nnz), current_frontier[:,0])) + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), current_frontier[:,0])) row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) - col_select_mtx_indices = torch.stack((next_frontier, torch.range(next_frontier.size(0))) + col_select_mtx_indices = torch.stack((next_frontier, torch.arange(start=0, end=next_frontier.size(0))\ + .cuda())) col_select_mtx_values = torch.cuda.FloatTensor(next_frontier.size(0)).fill_(1.0) # multiply row_select matrix with adj_matrix + sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), + row_select_mtx_values, + adj_matrix._indices(), adj_matrix._values(), + nnz, node_count, node_count, coalesced=True) # multiply adj_matrix with col_select matrix + sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, + col_select_mtx_indices.long(), col_select_mtx_values, + nnz, node_count, next_frontier.size(0), coalesced=True) # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] - current_frontier[:,0] = frontier_next + current_frontier[:,0] = next_frontier + adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ + size=(nnz, next_frontier.size(0))) + adj_matrices[i] = adj_matrix_sample - return adj_matrices + return batches[:, 0], next_frontier, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 57c759d..ac38b46 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -9,6 +9,7 @@ import torch.nn as nn import torch.autograd.profiler as profiler import torch_geometric +from torch_geometric.data import Data, Dataset from torch_geometric.datasets import Planetoid, Reddit import torch_geometric.transforms as T from torch_geometric.utils import add_remaining_self_loops @@ -130,21 +131,26 @@ def split_coo(adj_matrix, node_count, n_per_proc, dim): def row_normalize(mx): """Row-normalize sparse matrix""" + print("sum + float_power", flush=True) 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 = r_inv._values() + print("set inf", flush=True) r_inv_values[torch.isinf(r_inv_values)] = 0 + print("construct tensor", flush=True) 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()) + print("normalize spgemm", flush=True) mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values().float(), mx._indices(), mx._values().float(), r_mat_inv.size(0), r_mat_inv.size(1), mx.size(1), coalesced=True) + print("mx construct tensor", flush=True) mx = torch.sparse_coo_tensor(indices=mx_indices, values=mx_values, size=(r_mat_inv.size(0), mx.size(1))) return mx @@ -303,11 +309,17 @@ def main(args): if rank_c >= (size // args.replication): return + print("start partitioning", flush=True) features_loc, g_loc, ampbyp = one5d_partition(rank, size, inputs, adj_matrix, data, \ inputs, num_classes, args.replication, \ args.normalize) + print("end partitioning", flush=True) + + print("coalescing", flush=True) g_loc = g_loc.coalesce() + print("normalizing", flush=True) g_loc = row_normalize(g_loc) + print("done normalizing", flush=True) features_loc = features_loc.to(device) g_loc = g_loc.to(device) @@ -325,9 +337,11 @@ def main(args): train_nid = data.train_mask.nonzero().squeeze() test_nid = data.test_mask.nonzero().squeeze() - result = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, args.n_layers, train_nid) + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, \ + args.n_layers, train_nid) - return result # return here while testing sampling code + print(f"sample: {adj_matrices}") + return current_frontier, next_frontier, adj_matrices, g_loc # return here while testing sampling code # create GCN model torch.manual_seed(0) From 0068e9ee1fa5945d5a82cf4c357af0e7f15eba63 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 28 Jan 2022 11:56:36 -0800 Subject: [PATCH 018/178] Added timers to ladies sampler --- .gitignore | 1 + cagnet/samplers/ladies.py | 110 ++++++++++++++++++++++++++++++++++++-- examples/gcn_15d.py | 24 +++++---- 3 files changed, 123 insertions(+), 12 deletions(-) diff --git a/.gitignore b/.gitignore index 3c5cbdf..a9e42f4 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ gcn_distr_3d.py gcn.py reddit.py tests/ +*.qdrep diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 901a7aa..789c453 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -1,20 +1,49 @@ import torch import torch_sparse +from collections import defaultdict + +def start_time(timer): + timer.record() + +def stop_time(start_timer, stop_timer): + stop_timer.record() + torch.cuda.synchronize() + return start_timer.elapsed_time(stop_timer) def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, train_nodes): + 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) + + total_start_timer = torch.cuda.Event(enable_timing=True) + total_stop_timer = torch.cuda.Event(enable_timing=True) + + timing_dict = defaultdict(list) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-instantiations") # frontiers = torch.cuda.IntTensor(n_layers, batch_size, mb_count) batches = torch.cuda.IntTensor(batch_size, mb_count) # initially the minibatch current_frontier = torch.cuda.IntTensor(batch_size + frontier_size, mb_count) # frontiers = torch.cuda.IntTensor(n_layers - 1, batch_size + frontier_size, mb_count) adj_matrices = [None] * n_layers node_count = adj_matrix.size(0) + torch.cuda.nvtx.range_pop() + print(f"instantiations: {stop_time(start_timer, stop_timer)}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") torch.manual_seed(0) # Generate minibatch vertices for i in range(mb_count): idx = torch.randperm(train_nodes.size(0))[:batch_size] batches[:,i] = train_nodes[idx] + torch.cuda.nvtx.range_pop() + print(f"get-minibatch-vtxs: {stop_time(start_timer, stop_timer)}") + start_time(total_start_timer) for i in range(n_layers): if i == 0: nnz = batches[:, 0].size(0) # assumes mb_count == 1 @@ -23,50 +52,106 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr # A * Q^i # indices would change based on mb_count + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-gen-sparse-frontier") if i == 0: frontier_indices = torch.stack((torch.cuda.IntTensor(nnz).fill_(0), batches[:,0])) else: frontier_indices = torch.stack((torch.cuda.IntTensor(nnz).fill_(0), current_frontier[:,0])) frontier_values = torch.cuda.FloatTensor(nnz).fill_(1.0) adj_mat_squared = torch.pow(adj_matrix._values(), 2) + torch.cuda.nvtx.range_pop() + print(f"gen-sparse-frontier: {stop_time(start_timer, stop_timer)}") + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-probability-spgemm") p_num_indices, p_num_values = torch_sparse.spspmm(frontier_indices.long(), frontier_values, adj_matrix._indices().long(), adj_mat_squared, mb_count, node_count, node_count, coalesced=True) + torch.cuda.nvtx.range_pop() + print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-compute-p") p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count)) p_den = torch.sparse.sum(p_num, dim=1) # Assumes mb_count = 1 p_den = p_den._values().item() p = torch.div(p_num, p_den) + torch.cuda.nvtx.range_pop() + print(f"compute-p: {stop_time(start_timer, stop_timer)}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-pre-loop") max_prob = torch.max(p._values()) next_frontier = torch.cuda.LongTensor(frontier_size).fill_(node_count) sampled_count = 0 + torch.cuda.nvtx.range_pop() + print(f"pre-loop: {stop_time(start_timer, stop_timer)}") + iter_count = 0 + torch.cuda.nvtx.range_push("nvtx-sampling") while sampled_count < frontier_size: + iter_count += 1 + start_time(sample_start_timer) + torch.cuda.nvtx.range_push("nvtx-sampling-iter") # sample u.a.r. sampled_count vertices per minibatch from vertices in the minibatch's neighborhood # neighbored is vertices with nonzero probabilities - # TODO: What if a vertex is sampled more than once (idk if w/o replacement is possible) + # TODO: What if a vertex is sampled more than once? (idk if w/o replacement is possible) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-rand-vtxs") sampled_verts_ids = torch.randint(p._nnz(), (frontier_size,)) + torch.cuda.nvtx.range_pop() + timing_dict["rand_vtxs"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-p-loc") p_loc = p._values()[sampled_verts_ids] + torch.cuda.nvtx.range_pop() + timing_dict["p_loc"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-dart-throw") dart_throws = torch.cuda.FloatTensor(frontier_size).uniform_(to=max_prob) + torch.cuda.nvtx.range_pop() + timing_dict["dart_throw"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_misses = p_loc <= dart_throws dart_hits = p_loc > dart_throws hit_count = dart_hits.sum() hit_verts_ids = sampled_verts_ids[dart_hits] sampled_verts = p._indices()[1][hit_verts_ids] + torch.cuda.nvtx.range_pop() + timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-set-probs") next_frontier[dart_hits] = sampled_verts p._values()[hit_verts_ids] = 0.0 + torch.cuda.nvtx.range_pop() + timing_dict["set_probs"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-count-samples") sampled_count = torch.sum(next_frontier != node_count) - + torch.cuda.nvtx.range_pop() + timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() # nvtx-sampling-iter + timing_dict["sampling_iters"].append(stop_time(sample_start_timer, sample_stop_timer)) + torch.cuda.nvtx.range_pop() # nvtx-sampling + + # print(f"iter_count: {iter_count}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-construct-nextf") next_frontier = torch.cat((next_frontier, batches[:,0]), dim=0) + torch.cuda.nvtx.range_pop() + print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-mtxs") if i == 0: row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches[:,0])) else: @@ -76,21 +161,40 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr col_select_mtx_indices = torch.stack((next_frontier, torch.arange(start=0, end=next_frontier.size(0))\ .cuda())) col_select_mtx_values = torch.cuda.FloatTensor(next_frontier.size(0)).fill_(1.0) + torch.cuda.nvtx.range_pop() + print(f"select-mtxs: {stop_time(start_timer, stop_timer)}") # multiply row_select matrix with adj_matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), row_select_mtx_values, adj_matrix._indices(), adj_matrix._values(), nnz, node_count, node_count, coalesced=True) + torch.cuda.nvtx.range_pop() + print(f"row-select-spgemm: {stop_time(start_timer, stop_timer)}") + # multiply adj_matrix with col_select matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, col_select_mtx_indices.long(), col_select_mtx_values, nnz, node_count, next_frontier.size(0), coalesced=True) + torch.cuda.nvtx.range_pop() + print(f"col-select-spgemm: {stop_time(start_timer, stop_timer)}") # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-set-sample") current_frontier[:,0] = next_frontier adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ size=(nnz, next_frontier.size(0))) adj_matrices[i] = adj_matrix_sample - + torch.cuda.nvtx.range_pop() + print(f"set-sample: {stop_time(start_timer, stop_timer)}") + + print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") + for k, v in timing_dict.items(): + print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") + print(f"iter_count: {iter_count}") return batches[:, 0], next_frontier, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index ac38b46..3f5ed3e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -131,26 +131,21 @@ def split_coo(adj_matrix, node_count, n_per_proc, dim): def row_normalize(mx): """Row-normalize sparse matrix""" - print("sum + float_power", flush=True) 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 = r_inv._values() - print("set inf", flush=True) r_inv_values[torch.isinf(r_inv_values)] = 0 - print("construct tensor", flush=True) 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()) - print("normalize spgemm", flush=True) mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values().float(), mx._indices(), mx._values().float(), r_mat_inv.size(0), r_mat_inv.size(1), mx.size(1), coalesced=True) - print("mx construct tensor", flush=True) mx = torch.sparse_coo_tensor(indices=mx_indices, values=mx_values, size=(r_mat_inv.size(0), mx.size(1))) return mx @@ -278,6 +273,7 @@ def main(args): 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 @@ -315,16 +311,17 @@ def main(args): args.normalize) print("end 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) + print("coalescing", flush=True) g_loc = g_loc.coalesce() print("normalizing", flush=True) g_loc = row_normalize(g_loc) print("done normalizing", 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) n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) @@ -337,8 +334,17 @@ def main(args): train_nid = data.train_mask.nonzero().squeeze() test_nid = data.test_mask.nonzero().squeeze() + # do it once before timing + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, \ + args.n_layers, train_nid) + + print() + torch.cuda.profiler.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("nvtx-sampler") current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, \ args.n_layers, train_nid) + torch.cuda.nvtx.range_pop() + torch.cuda.profiler.cudart().cudaProfilerStop() print(f"sample: {adj_matrices}") return current_frontier, next_frontier, adj_matrices, g_loc # return here while testing sampling code From 715710b264945ab7f13c4a3ca20c77d5fcdb71bf Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 2 Feb 2022 18:29:04 -0800 Subject: [PATCH 019/178] Shifting code to bulk version --- cagnet/samplers/ladies.py | 39 +++++++++++++++++++++++++-------------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 789c453..0e5c8a6 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -25,8 +25,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-instantiations") # frontiers = torch.cuda.IntTensor(n_layers, batch_size, mb_count) - batches = torch.cuda.IntTensor(batch_size, mb_count) # initially the minibatch - current_frontier = torch.cuda.IntTensor(batch_size + frontier_size, mb_count) + batches = torch.cuda.IntTensor(mb_count, batch_size) # initially the minibatch, note row-major + current_frontier = torch.cuda.IntTensor(mb_count, batch_size + frontier_size) # frontiers = torch.cuda.IntTensor(n_layers - 1, batch_size + frontier_size, mb_count) adj_matrices = [None] * n_layers node_count = adj_matrix.size(0) @@ -39,25 +39,31 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr # Generate minibatch vertices for i in range(mb_count): idx = torch.randperm(train_nodes.size(0))[:batch_size] - batches[:,i] = train_nodes[idx] + batches[i,:] = train_nodes[idx] torch.cuda.nvtx.range_pop() print(f"get-minibatch-vtxs: {stop_time(start_timer, stop_timer)}") start_time(total_start_timer) for i in range(n_layers): if i == 0: - nnz = batches[:, 0].size(0) # assumes mb_count == 1 + nnz = batches[0, :].size(0) else: - nnz = current_frontier[:, 0].size(0) # assumes mb_count == 1 + nnz = current_frontier[0, :].size(0) # A * Q^i # indices would change based on mb_count start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-sparse-frontier") if i == 0: - frontier_indices = torch.stack((torch.cuda.IntTensor(nnz).fill_(0), batches[:,0])) + frontier_indices_rows = torch.arange(mb_count).cuda() + frontier_indices_rows = frontier_indices_rows.repeat_interleave(nnz) + frontier_indices_cols = batches.view(-1) + frontier_indices = torch.stack((frontier_indices_rows, frontier_indices_cols)) else: - frontier_indices = torch.stack((torch.cuda.IntTensor(nnz).fill_(0), current_frontier[:,0])) + frontier_indices_rows = torch.arange(mb_count).cuda() + frontier_indices_rows = frontier_indices_rows.repeat_interleave(nnz) + forntier_indices_cols = current_frontier.view(-1) + frontier_indices = torch.stack((frontier_indices_rows, frontier_indices_cols)) frontier_values = torch.cuda.FloatTensor(nnz).fill_(1.0) adj_mat_squared = torch.pow(adj_matrix._values(), 2) torch.cuda.nvtx.range_pop() @@ -71,6 +77,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") + print(f"p_num_indices.size(): {p_num_indices.size()}") + print(f"p_num_values.size(): {p_num_values.size()}") + start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count)) @@ -84,15 +93,17 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") - max_prob = torch.max(p._values()) - next_frontier = torch.cuda.LongTensor(frontier_size).fill_(node_count) - sampled_count = 0 + # max_prob = torch.sparse.max(p, dim=1) + p_torch_sparse = torch_sparse.tensor.SparseTensor.from_torch_sparse_coo_tensor(p) + max_prob = torch_sparse.reduce.reduction(p_torch_sparse, dim=1, reduce="max") + next_frontier = torch.cuda.LongTensor(mb_count, frontier_size).fill_(node_count) + sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) torch.cuda.nvtx.range_pop() print(f"pre-loop: {stop_time(start_timer, stop_timer)}") - + iter_count = 0 torch.cuda.nvtx.range_push("nvtx-sampling") - while sampled_count < frontier_size: + while (sampled_count < frontier_size).all(): iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -101,13 +112,13 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr # TODO: What if a vertex is sampled more than once? (idk if w/o replacement is possible) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-rand-vtxs") - sampled_verts_ids = torch.randint(p._nnz(), (frontier_size,)) + sampled_verts_ids = torch.randint(p._nnz(), (mb_count, frontier_size)) torch.cuda.nvtx.range_pop() timing_dict["rand_vtxs"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-p-loc") - p_loc = p._values()[sampled_verts_ids] + p_loc = p._values()[sampled_verts_ids] torch.cuda.nvtx.range_pop() timing_dict["p_loc"].append(stop_time(start_timer, stop_timer)) From 531d6e5431e7b5fba2696f439544d6ca1d4b16b1 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 3 Feb 2022 15:23:58 -0800 Subject: [PATCH 020/178] Added nccl p2p test --- cagnet/samplers/ladies.py | 100 +++++++++++++++++++++----------------- test_nccl_p2p.py | 38 +++++++++++++++ 2 files changed, 94 insertions(+), 44 deletions(-) create mode 100644 test_nccl_p2p.py diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0e5c8a6..6fbd70f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -112,19 +112,29 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr # TODO: What if a vertex is sampled more than once? (idk if w/o replacement is possible) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-rand-vtxs") - sampled_verts_ids = torch.randint(p._nnz(), (mb_count, frontier_size)) + sampled_verts_ids = torch.cuda.LongTensor(mb_count, frontier_size) # indices into p._values() + nnz_count = 0 + for j in range(mb_count): + nnz_mb_mask = p._indices()[0, :] == j + p_values_mb = p._values()[nnz_mb_mask] + sampled_verts_ids[i] = torch.randint(p_values_mb.size(0), (frontier_size,)) + sampled_verts_ids[i] += nnz_count + nnz_count += p_values_mb.size(0) torch.cuda.nvtx.range_pop() timing_dict["rand_vtxs"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-p-loc") - p_loc = p._values()[sampled_verts_ids] + p_loc = p._values()[sampled_verts_ids.view(-1)].view(mb_count, frontier_size) + torch.cuda.nvtx.range_pop() timing_dict["p_loc"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - dart_throws = torch.cuda.FloatTensor(frontier_size).uniform_(to=max_prob) + dart_throws = torch.cuda.FloatTensor(mb_count, frontier_size) + for j in range(mb_count): + dart_throws[i] = torch.cuda.FloatTensor(frontier_size).uniform_(to=max_prob[i]) torch.cuda.nvtx.range_pop() timing_dict["dart_throw"].append(stop_time(start_timer, stop_timer)) @@ -132,7 +142,6 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_misses = p_loc <= dart_throws dart_hits = p_loc > dart_throws - hit_count = dart_hits.sum() hit_verts_ids = sampled_verts_ids[dart_hits] sampled_verts = p._indices()[1][hit_verts_ids] torch.cuda.nvtx.range_pop() @@ -157,52 +166,55 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr # print(f"iter_count: {iter_count}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-construct-nextf") - next_frontier = torch.cat((next_frontier, batches[:,0]), dim=0) + next_frontier = torch.cat((next_frontier, batches), dim=1) torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-mtxs") - if i == 0: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches[:,0])) - else: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), current_frontier[:,0])) - row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) - - col_select_mtx_indices = torch.stack((next_frontier, torch.arange(start=0, end=next_frontier.size(0))\ - .cuda())) - col_select_mtx_values = torch.cuda.FloatTensor(next_frontier.size(0)).fill_(1.0) - torch.cuda.nvtx.range_pop() - print(f"select-mtxs: {stop_time(start_timer, stop_timer)}") + for j in range(mb_count): + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-mtxs") + if i == 0: + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches[j,:])) + else: + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), \ + current_frontier[j, :])) + row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) + + col_select_mtx_indices = torch.stack((next_frontier[j], torch.arange(start=0, \ + end=next_frontier[j].size(0)).cuda())) + col_select_mtx_values = torch.cuda.FloatTensor(next_frontier[j].size(0)).fill_(1.0) + torch.cuda.nvtx.range_pop() + print(f"select-mtxs: {stop_time(start_timer, stop_timer)}") - # multiply row_select matrix with adj_matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), - row_select_mtx_values, - adj_matrix._indices(), adj_matrix._values(), - nnz, node_count, node_count, coalesced=True) - torch.cuda.nvtx.range_pop() - print(f"row-select-spgemm: {stop_time(start_timer, stop_timer)}") + # multiply row_select matrix with adj_matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") + sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), + row_select_mtx_values, + adj_matrix._indices(), adj_matrix._values(), + nnz, node_count, node_count, coalesced=True) + torch.cuda.nvtx.range_pop() + print(f"row-select-spgemm: {stop_time(start_timer, stop_timer)}") - # multiply adj_matrix with col_select matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, - col_select_mtx_indices.long(), col_select_mtx_values, - nnz, node_count, next_frontier.size(0), coalesced=True) - torch.cuda.nvtx.range_pop() - print(f"col-select-spgemm: {stop_time(start_timer, stop_timer)}") - # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] + # multiply adj_matrix with col_select matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") + sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, + col_select_mtx_indices.long(), col_select_mtx_values, + nnz, node_count, next_frontier[j].size(0), + coalesced=True) + torch.cuda.nvtx.range_pop() + print(f"col-select-spgemm: {stop_time(start_timer, stop_timer)}") + # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-sample") - current_frontier[:,0] = next_frontier - adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ - size=(nnz, next_frontier.size(0))) - adj_matrices[i] = adj_matrix_sample - torch.cuda.nvtx.range_pop() - print(f"set-sample: {stop_time(start_timer, stop_timer)}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-set-sample") + current_frontier[j, :] = next_frontier[j, :] + adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ + size=(nnz, next_frontier[j].size(0))) + adj_matrices[i] = adj_matrix_sample + torch.cuda.nvtx.range_pop() + print(f"set-sample: {stop_time(start_timer, stop_timer)}") print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") for k, v in timing_dict.items(): diff --git a/test_nccl_p2p.py b/test_nccl_p2p.py new file mode 100644 index 0000000..a6030ae --- /dev/null +++ b/test_nccl_p2p.py @@ -0,0 +1,38 @@ +import argparse +import os +import torch +import torch.distributed as dist +import socket + +def test_nccl(args): + 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="nccl") + rank = dist.get_rank() + size = dist.get_world_size() + print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}") + print(f"backend: {dist.get_backend()}") + + if rank == 0: + x = torch.cuda.IntTensor(5).fill_(3) # vector of size 5, set to all 3's + dist.send(x, dst=1) + elif rank == 1: + y = torch.cuda.IntTensor(5).fill_(2) # vector of size 5, set to all 2's + print(f"y_before: {y}") + dist.recv(y, src=0) + print(f"y_after: {y}") + +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') + + args = parser.parse_args() + test_nccl(args) From cde046b871e6881ba13957829df398a6f85ce406 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 4 Feb 2022 15:58:10 -0800 Subject: [PATCH 021/178] Fixing bulk implementation --- cagnet/samplers/ladies.py | 10 ++++++---- examples/gcn_15d.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 6fbd70f..2ea4503 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -64,7 +64,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr frontier_indices_rows = frontier_indices_rows.repeat_interleave(nnz) forntier_indices_cols = current_frontier.view(-1) frontier_indices = torch.stack((frontier_indices_rows, frontier_indices_cols)) - frontier_values = torch.cuda.FloatTensor(nnz).fill_(1.0) + frontier_values = torch.cuda.FloatTensor(nnz * mb_count).fill_(1.0) adj_mat_squared = torch.pow(adj_matrix._values(), 2) torch.cuda.nvtx.range_pop() print(f"gen-sparse-frontier: {stop_time(start_timer, stop_timer)}") @@ -85,9 +85,11 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count)) p_den = torch.sparse.sum(p_num, dim=1) - # Assumes mb_count = 1 - p_den = p_den._values().item() - p = torch.div(p_num, p_den) + for j in range(mb_count): + if p_den._nnz() != mb_count: + print("ERROR nnz: {p_den._nnz()} mb_count: {mb_count}") + p_den_mb = p_den._values()[j].item() + p = torch.div(p_num, p_den) torch.cuda.nvtx.range_pop() print(f"compute-p: {stop_time(start_timer, stop_timer)}") diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 3f5ed3e..ed79618 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -335,7 +335,7 @@ def main(args): test_nid = data.test_mask.nonzero().squeeze() # do it once before timing - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, \ + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 2, \ args.n_layers, train_nid) print() From 7f2a7e83429023792da7b23efdd44efbec724a37 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 8 Feb 2022 20:08:40 -0800 Subject: [PATCH 022/178] Initial bulk LADIES implementation (mb_count=2) --- cagnet/samplers/ladies.py | 17 +++++++++-------- examples/gcn_15d.py | 12 ++++++++---- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 2ea4503..8d8c641 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -28,7 +28,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr batches = torch.cuda.IntTensor(mb_count, batch_size) # initially the minibatch, note row-major current_frontier = torch.cuda.IntTensor(mb_count, batch_size + frontier_size) # frontiers = torch.cuda.IntTensor(n_layers - 1, batch_size + frontier_size, mb_count) - adj_matrices = [None] * n_layers + adj_matrices = [[None] * n_layers for x in range(mb_count)] # adj_matrices[i][j] -- mb i layer j node_count = adj_matrix.size(0) torch.cuda.nvtx.range_pop() print(f"instantiations: {stop_time(start_timer, stop_timer)}") @@ -82,14 +82,15 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count)) + p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ + size=(mb_count, node_count)) p_den = torch.sparse.sum(p_num, dim=1) for j in range(mb_count): if p_den._nnz() != mb_count: print("ERROR nnz: {p_den._nnz()} mb_count: {mb_count}") p_den_mb = p_den._values()[j].item() - p = torch.div(p_num, p_den) + p = torch.div(p_num, p_den_mb) torch.cuda.nvtx.range_pop() print(f"compute-p: {stop_time(start_timer, stop_timer)}") @@ -105,7 +106,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr iter_count = 0 torch.cuda.nvtx.range_push("nvtx-sampling") - while (sampled_count < frontier_size).all(): + while (sampled_count < mb_count * frontier_size).all(): iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -119,8 +120,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr for j in range(mb_count): nnz_mb_mask = p._indices()[0, :] == j p_values_mb = p._values()[nnz_mb_mask] - sampled_verts_ids[i] = torch.randint(p_values_mb.size(0), (frontier_size,)) - sampled_verts_ids[i] += nnz_count + sampled_verts_ids[j] = torch.randint(p_values_mb.size(0), (frontier_size,)) + sampled_verts_ids[j] += nnz_count nnz_count += p_values_mb.size(0) torch.cuda.nvtx.range_pop() timing_dict["rand_vtxs"].append(stop_time(start_timer, stop_timer)) @@ -214,7 +215,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr current_frontier[j, :] = next_frontier[j, :] adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ size=(nnz, next_frontier[j].size(0))) - adj_matrices[i] = adj_matrix_sample + adj_matrices[j][i] = adj_matrix_sample torch.cuda.nvtx.range_pop() print(f"set-sample: {stop_time(start_timer, stop_timer)}") @@ -222,4 +223,4 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr for k, v in timing_dict.items(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") print(f"iter_count: {iter_count}") - return batches[:, 0], next_frontier, adj_matrices + return batches, next_frontier, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index ed79618..cbe6dca 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -335,14 +335,16 @@ def main(args): test_nid = data.test_mask.nonzero().squeeze() # do it once before timing - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 2, \ - args.n_layers, train_nid) + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, \ + args.n_bulkmb, args.n_layers, \ + train_nid) print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, 1, \ - args.n_layers, train_nid) + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, \ + args.n_bulkmb, args.n_layers, \ + train_nid) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() @@ -462,6 +464,8 @@ def main(args): 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') args = parser.parse_args() print(args) From f5eccc5db727e69ac1cceae801fd44860348ee31 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 8 Feb 2022 20:11:40 -0800 Subject: [PATCH 023/178] Added LADIES test file --- .gitignore | 5 +- tests/test_ladies.py | 166 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 1 deletion(-) create mode 100644 tests/test_ladies.py diff --git a/.gitignore b/.gitignore index a9e42f4..7ecaa69 100644 --- a/.gitignore +++ b/.gitignore @@ -20,5 +20,8 @@ gcn_distr_2d.py gcn_distr_3d.py gcn.py reddit.py -tests/ +tests/data +tests/examples +tests/LADIES/ +tests/*.out *.qdrep diff --git a/tests/test_ladies.py b/tests/test_ladies.py new file mode 100644 index 0000000..2c64fc2 --- /dev/null +++ b/tests/test_ladies.py @@ -0,0 +1,166 @@ +import argparse +import examples.gcn_15d +import LADIES.pytorch_ladies + +import numpy as np +import torch + +def main(args): + """ + # Testing sampled neighbors, probability distribution, and normalized Laplacian + # with original LADIES code + + mine, mine_mtx, mine_batch = examples.gcn_15d.main(args) + mine = mine.cuda() + them, them_mtx, them_batch = LADIES.pytorch_ladies.main(args) + them = torch.from_numpy(them).cuda() + + mine_sampled_neighbors = mine._indices()[1,:].squeeze() + them_sampled_neighbors = np.nonzero(them).squeeze() + + mine_sampled_probs = mine._values().float() + them_sampled_probs = them[them_sampled_neighbors].float() + + if not torch.equal(mine_sampled_neighbors, them_sampled_neighbors): + print("indices don't match") + print(f"mine_sampled_neighbors.size: {mine_sampled_neighbors.size()}") + print(f"mine_sampled_neighbors: {mine_sampled_neighbors}") + print(f"them_sampled_neighbors.size: {them_sampled_neighbors.size()}") + print(f"them_sampled_neighbors: {them_sampled_neighbors}") + else: + print("indices match") + + if not torch.allclose(mine_sampled_probs, them_sampled_probs, atol=0.001, rtol=0.0): + print("probabilities don't match") + print(f"mine_sampled_probs.size: {mine_sampled_probs.size()}") + print(f"mine_sampled_probs: {mine_sampled_probs}") + print(f"them_sampled_probs.size: {them_sampled_probs.size()}") + print(f"them_sampled_probs: {them_sampled_probs}") + else: + print("probabilities match") + + mine_mtx_normalization = mine_mtx._values() + them_mtx_normalization = torch.from_numpy(them_mtx.data).float().cuda() + + if not torch.allclose(mine_mtx_normalization, them_mtx_normalization, atol=0.001, rtol=0.0): + print("laplacians don't match") + unmatch_idx = ~torch.isclose(mine_mtx_normalization, them_mtx_normalization, atol=0.001, rtol=0.0) + print(f"unmatch_idx: {unmatch_idx.nonzero()}") + print(f"mine_mtx_normalization: {mine_mtx_normalization[unmatch_idx]}") + print(f"them_mtx_normalization: {them_mtx_normalization[unmatch_idx]}") + else: + print("laplacians match") + + them_batch = torch.from_numpy(them_batch).cuda() + + if not torch.equal(mine_batch, them_batch): + print("batches don't match") + print(f"mine_batch: {mine_batch}") + print(f"them_batch: {them_batch}") + else: + print("batches match") + """ + + + """ + 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") + for i in range(args.n_bulkmb): + current_frontier = current_frontier_all[i] # current_frontier for minibatch i + next_frontier = next_frontier_all[i] # 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 + 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 adj_matrix._values()[nnz_count] != neighbors_nnz[v_vtx_idx]: + 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') + args = parser.parse_args() + print(args) + + main(args) From b30b09c01ba18ced2721f5068ac0efa7a8c23e46 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 10 Feb 2022 16:49:03 -0800 Subject: [PATCH 024/178] Recomputes ceiling for sampling per iteration --- cagnet/samplers/ladies.py | 26 +++++++++++++++++--------- tests/test_ladies.py | 1 + 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 8d8c641..d5ec7fc 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -62,7 +62,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr else: frontier_indices_rows = torch.arange(mb_count).cuda() frontier_indices_rows = frontier_indices_rows.repeat_interleave(nnz) - forntier_indices_cols = current_frontier.view(-1) + frontier_indices_cols = current_frontier.view(-1) frontier_indices = torch.stack((frontier_indices_rows, frontier_indices_cols)) frontier_values = torch.cuda.FloatTensor(nnz * mb_count).fill_(1.0) adj_mat_squared = torch.pow(adj_matrix._values(), 2) @@ -96,9 +96,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") - # max_prob = torch.sparse.max(p, dim=1) - p_torch_sparse = torch_sparse.tensor.SparseTensor.from_torch_sparse_coo_tensor(p) - max_prob = torch_sparse.reduce.reduction(p_torch_sparse, dim=1, reduce="max") + # p_torch_sparse = torch_sparse.tensor.SparseTensor.from_torch_sparse_coo_tensor(p) + # max_prob = torch_sparse.reduce.reduction(p_torch_sparse, dim=1, reduce="max") next_frontier = torch.cuda.LongTensor(mb_count, frontier_size).fill_(node_count) sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) torch.cuda.nvtx.range_pop() @@ -110,7 +109,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") - # sample u.a.r. sampled_count vertices per minibatch from vertices in the minibatch's neighborhood + # sample u.a.r. sampled_count vertices per minibatch from vertices in the mb's neighborhood # neighbored is vertices with nonzero probabilities # TODO: What if a vertex is sampled more than once? (idk if w/o replacement is possible) start_time(start_timer) @@ -126,6 +125,13 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr torch.cuda.nvtx.range_pop() timing_dict["rand_vtxs"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-comp-maxprob") + p_torch_sparse = torch_sparse.tensor.SparseTensor.from_torch_sparse_coo_tensor(p) + max_prob = torch_sparse.reduce.reduction(p_torch_sparse, dim=1, reduce="max") + torch.cuda.nvtx.range_pop() + timing_dict["max-prob"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-p-loc") p_loc = p._values()[sampled_verts_ids.view(-1)].view(mb_count, frontier_size) @@ -173,6 +179,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") + torch.cuda.nvtx.range_push("nvtx-select-rowcols") for j in range(mb_count): start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-mtxs") @@ -187,7 +194,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr end=next_frontier[j].size(0)).cuda())) col_select_mtx_values = torch.cuda.FloatTensor(next_frontier[j].size(0)).fill_(1.0) torch.cuda.nvtx.range_pop() - print(f"select-mtxs: {stop_time(start_timer, stop_timer)}") + timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) # multiply row_select matrix with adj_matrix start_time(start_timer) @@ -197,7 +204,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr adj_matrix._indices(), adj_matrix._values(), nnz, node_count, node_count, coalesced=True) torch.cuda.nvtx.range_pop() - print(f"row-select-spgemm: {stop_time(start_timer, stop_timer)}") + timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) # multiply adj_matrix with col_select matrix start_time(start_timer) @@ -207,7 +214,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr nnz, node_count, next_frontier[j].size(0), coalesced=True) torch.cuda.nvtx.range_pop() - print(f"col-select-spgemm: {stop_time(start_timer, stop_timer)}") + timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] start_time(start_timer) @@ -217,7 +224,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr size=(nnz, next_frontier[j].size(0))) adj_matrices[j][i] = adj_matrix_sample torch.cuda.nvtx.range_pop() - print(f"set-sample: {stop_time(start_timer, stop_timer)}") + timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") for k, v in timing_dict.items(): diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 2c64fc2..e695832 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -74,6 +74,7 @@ def main(args): print("") print("Beginning LADIES sampler test") for i in range(args.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 adj_matrices = adj_matrices_all[i] # adj_matrices for minibatch i From 4a8f681d3fef1a2502ab71ecef7c1731a9697d1b Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 15 Feb 2022 16:31:28 -0800 Subject: [PATCH 025/178] Generate distinct minibatches --- cagnet/samplers/ladies.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index d5ec7fc..14bb0a7 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -35,10 +35,12 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") + torch.manual_seed(0) + vertex_perm = torch.randperm(train_nodes.size(0)) # Generate minibatch vertices for i in range(mb_count): - idx = torch.randperm(train_nodes.size(0))[:batch_size] + idx = vertex_perm[(i * batch_size):((i + 1) * batch_size)] batches[i,:] = train_nodes[idx] torch.cuda.nvtx.range_pop() print(f"get-minibatch-vtxs: {stop_time(start_timer, stop_timer)}") From 5a1b5760887792be694d708572369d9010dab1a2 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 16 Feb 2022 13:23:10 -0800 Subject: [PATCH 026/178] Implementing dart throwing for all neighborhood vtxs --- cagnet/samplers/ladies.py | 63 +++++++++++++++++++++------------------ 1 file changed, 34 insertions(+), 29 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 14bb0a7..1fb9f26 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -111,21 +111,6 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") - # sample u.a.r. sampled_count vertices per minibatch from vertices in the mb's neighborhood - # neighbored is vertices with nonzero probabilities - # TODO: What if a vertex is sampled more than once? (idk if w/o replacement is possible) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-rand-vtxs") - sampled_verts_ids = torch.cuda.LongTensor(mb_count, frontier_size) # indices into p._values() - nnz_count = 0 - for j in range(mb_count): - nnz_mb_mask = p._indices()[0, :] == j - p_values_mb = p._values()[nnz_mb_mask] - sampled_verts_ids[j] = torch.randint(p_values_mb.size(0), (frontier_size,)) - sampled_verts_ids[j] += nnz_count - nnz_count += p_values_mb.size(0) - torch.cuda.nvtx.range_pop() - timing_dict["rand_vtxs"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-comp-maxprob") @@ -134,29 +119,49 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr torch.cuda.nvtx.range_pop() timing_dict["max-prob"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-p-loc") - p_loc = p._values()[sampled_verts_ids.view(-1)].view(mb_count, frontier_size) - - torch.cuda.nvtx.range_pop() - timing_dict["p_loc"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - dart_throws = torch.cuda.FloatTensor(mb_count, frontier_size) + dart_throws_values = torch.cuda.FloatTensor(p._nnz()) + mb_start = 0 for j in range(mb_count): - dart_throws[i] = torch.cuda.FloatTensor(frontier_size).uniform_(to=max_prob[i]) + neighbors_count = (p._indices()[0,:] == j).sum().item() + mb_end = mb_start + neighbors_count + dart_throws_values[mb_start:mb_end] = \ + torch.cuda.FloatTensor(neighbors_count).uniform_(to=max_prob[j]) + mb_start = mb_end torch.cuda.nvtx.range_pop() - timing_dict["dart_throw"].append(stop_time(start_timer, stop_timer)) + timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") - dart_misses = p_loc <= dart_throws - dart_hits = p_loc > dart_throws - hit_verts_ids = sampled_verts_ids[dart_hits] - sampled_verts = p._indices()[1][hit_verts_ids] + print(f"p._values.size(): {p._values().size()}") + print(f"p._values: {p._values()}") + print(f"dart_throws_values.size(): {dart_throws_values.size()}") + print(f"dart_throws_values: {dart_throws_values}") + dart_hits_values = p._values() >= dart_throws_values + print(f"max_prob: {max_prob}") + print(f"dart_hits_values.sum(): {dart_hits_values.sum()}") + print(f"dart_hits_values: {dart_hits_values}") torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-filter-hits") + for j in range(mb_count): + neighbors_mask = p._indices()[0,:] == j + neighbors = (p._indices()[1,:])[neighbors_mask] + neighbors_count = neighbors_mask.sum() + + hit_vtxs = torch.randperm(neighbors_count) + hit_vtxs = hit_vtxs[:frontier_size] + hit_vtxs = neighbors[hit_vtxs] + torch.cuda.nvtx.range_pop() + timing_dict["filter-hits"].append(stop_time(start_timer, stop_timer)) + + for k, v in timing_dict.items(): + print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") + print(f"iter_count: {iter_count}") + return None, None, None start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") From e2337d1df90e451f5ae73c65b0a828bac478a8c7 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 16 Feb 2022 18:32:24 -0800 Subject: [PATCH 027/178] Refactored to use new sampling algorithm --- cagnet/samplers/ladies.py | 71 +++++++++++++++++---------------------- 1 file changed, 30 insertions(+), 41 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 1fb9f26..60a275f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -98,16 +98,17 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") - # p_torch_sparse = torch_sparse.tensor.SparseTensor.from_torch_sparse_coo_tensor(p) - # max_prob = torch_sparse.reduce.reduction(p_torch_sparse, dim=1, reduce="max") - next_frontier = torch.cuda.LongTensor(mb_count, frontier_size).fill_(node_count) + next_frontier = torch.sparse_coo_tensor(indices=p._indices(), + values=torch.cuda.LongTensor(p._nnz()).fill_(0), + size=(mb_count, node_count)) sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) torch.cuda.nvtx.range_pop() print(f"pre-loop: {stop_time(start_timer, stop_timer)}") iter_count = 0 + print(f"p: {p}") torch.cuda.nvtx.range_push("nvtx-sampling") - while (sampled_count < mb_count * frontier_size).all(): + while (sampled_count < frontier_size).all(): iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -119,69 +120,57 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr torch.cuda.nvtx.range_pop() timing_dict["max-prob"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-gen-darts") + dart_throws_values = torch.cuda.FloatTensor(p._nnz()).uniform_() + torch.cuda.nvtx.range_pop() + timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - dart_throws_values = torch.cuda.FloatTensor(p._nnz()) - mb_start = 0 - for j in range(mb_count): - neighbors_count = (p._indices()[0,:] == j).sum().item() - mb_end = mb_start + neighbors_count - dart_throws_values[mb_start:mb_end] = \ - torch.cuda.FloatTensor(neighbors_count).uniform_(to=max_prob[j]) - mb_start = mb_end + dart_throws_values = dart_throws_values * torch.gather(max_prob, 0, p._indices()[0, :]) torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") - print(f"p._values.size(): {p._values().size()}") - print(f"p._values: {p._values()}") - print(f"dart_throws_values.size(): {dart_throws_values.size()}") - print(f"dart_throws_values: {dart_throws_values}") - dart_hits_values = p._values() >= dart_throws_values - print(f"max_prob: {max_prob}") - print(f"dart_hits_values.sum(): {dart_hits_values.sum()}") - print(f"dart_hits_values: {dart_hits_values}") + dart_hits_mask = dart_throws_values < p._values() torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-filter-hits") - for j in range(mb_count): - neighbors_mask = p._indices()[0,:] == j - neighbors = (p._indices()[1,:])[neighbors_mask] - neighbors_count = neighbors_mask.sum() - - hit_vtxs = torch.randperm(neighbors_count) - hit_vtxs = hit_vtxs[:frontier_size] - hit_vtxs = neighbors[hit_vtxs] + torch.cuda.nvtx.range_push("nvtx-set-probs") + p._values()[dart_hits_mask] = 0.0 torch.cuda.nvtx.range_pop() - timing_dict["filter-hits"].append(stop_time(start_timer, stop_timer)) + timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) - for k, v in timing_dict.items(): - print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") - print(f"iter_count: {iter_count}") - return None, None, None - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-probs") - next_frontier[dart_hits] = sampled_verts - p._values()[hit_verts_ids] = 0.0 + torch.cuda.nvtx.range_push("nvtx-add-to-frontier") + next_frontier_values = torch.logical_or(dart_hits_mask.int(), next_frontier._values()).long() + next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=next_frontier_values, + size=(mb_count, node_count)) torch.cuda.nvtx.range_pop() - timing_dict["set_probs"].append(stop_time(start_timer, stop_timer)) + timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-count-samples") - sampled_count = torch.sum(next_frontier != node_count) + sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) torch.cuda.nvtx.range_pop() # nvtx-sampling-iter timing_dict["sampling_iters"].append(stop_time(sample_start_timer, sample_stop_timer)) torch.cuda.nvtx.range_pop() # nvtx-sampling - # print(f"iter_count: {iter_count}") + # TODO: Might need to downsample if > frontier_size vertices were sampled for a minibatch + start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-construct-nextf") + print(f"before next_frontier: {next_frontier}") + next_frontier = torch.masked_select(next_frontier._indices()[1,:], \ + next_frontier._values().bool()).view(mb_count, frontier_size) + print(f"batches: {batches}") + print(f"after next_frontier: {next_frontier}") next_frontier = torch.cat((next_frontier, batches), dim=1) torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") From 331b8496eb8acd1d40d922d7ec52a841c69d3760 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 24 Feb 2022 20:13:13 -0800 Subject: [PATCH 028/178] Added downsampling at end of sampling loop --- cagnet/samplers/ladies.py | 18 ++++++++---- sparse-extension/setup.py | 2 +- ...se_coo_tensor.cpp => sparse_coo_tensor.cu} | 28 +++++++++++++++++++ 3 files changed, 42 insertions(+), 6 deletions(-) rename sparse-extension/{sparse_coo_tensor.cpp => sparse_coo_tensor.cu} (89%) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 60a275f..6781388 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -2,6 +2,8 @@ import torch_sparse from collections import defaultdict +from sparse_coo_tensor_cpp import downsample_gpu + def start_time(timer): timer.record() @@ -106,9 +108,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr print(f"pre-loop: {stop_time(start_timer, stop_timer)}") iter_count = 0 - print(f"p: {p}") torch.cuda.nvtx.range_push("nvtx-sampling") - while (sampled_count < frontier_size).all(): + while (sampled_count < frontier_size).any(): iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -158,19 +159,26 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-downsample") + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + downsample_gpu(next_frontier._values(), next_frontier._indices()[0,:], overflow, next_frontier._nnz()) + torch.cuda.nvtx.range_pop() + timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() # nvtx-sampling-iter timing_dict["sampling_iters"].append(stop_time(sample_start_timer, sample_stop_timer)) + torch.cuda.nvtx.range_pop() # nvtx-sampling # TODO: Might need to downsample if > frontier_size vertices were sampled for a minibatch + print(f"next_frontier.frontier-sizes: {torch.sparse.sum(next_frontier, dim=1)}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-construct-nextf") - print(f"before next_frontier: {next_frontier}") next_frontier = torch.masked_select(next_frontier._indices()[1,:], \ next_frontier._values().bool()).view(mb_count, frontier_size) - print(f"batches: {batches}") - print(f"after next_frontier: {next_frontier}") next_frontier = torch.cat((next_frontier, batches), dim=1) torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index e81e5fe..68f5f7b 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'], + ext_modules=[cpp_extension.CUDAExtension('sparse_coo_tensor_cpp', ['sparse_coo_tensor.cu'], extra_compile_args=["-lcusparse"])], cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/sparse-extension/sparse_coo_tensor.cpp b/sparse-extension/sparse_coo_tensor.cu similarity index 89% rename from sparse-extension/sparse_coo_tensor.cpp rename to sparse-extension/sparse_coo_tensor.cu index e6448b1..c0f495a 100644 --- a/sparse-extension/sparse_coo_tensor.cpp +++ b/sparse-extension/sparse_coo_tensor.cu @@ -226,7 +226,35 @@ void spmm_gpu(const at::Tensor& A_rowindices, C.t_(); } +__global__ void DownSample(long *q_values, long *q_rows, int *overflow, int nnz) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < nnz; i += stride) { + long row = q_rows[i]; + if (q_values[i] == 1 && ((int) atomicSub((unsigned int *)&overflow[row], 1)) > 0) { + q_values[i] = 0; + } + } +} + +void downsample_gpu(const at::Tensor& q_values, + const at::Tensor& q_rows, + 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<<>>(q_values.data(), q_rows.data(), + overflow.data(), nnz); + CHECK_ERROR("downsampling error") +} + 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"); + m.def("downsample_gpu", &downsample_gpu, "Downsampling for LADIES sampling algorithm"); } From 15c9218a85d0a38eb661d7f8a53c2adac3dc95b1 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 8 Mar 2022 14:57:16 -0800 Subject: [PATCH 029/178] Added 2D dart throwing --- cagnet/samplers/ladies.py | 24 +++++++-- examples/gcn_15d.py | 6 ++- sparse-extension/sparse_coo_tensor.cu | 76 +++++++++++++++++++++++++++ tests/test_ladies.py | 2 + 4 files changed, 101 insertions(+), 7 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 6781388..0f46910 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -2,7 +2,7 @@ import torch_sparse from collections import defaultdict -from sparse_coo_tensor_cpp import downsample_gpu +from sparse_coo_tensor_cpp import downsample_gpu, compute_darts_gpu, throw_darts_gpu def start_time(timer): timer.record() @@ -12,7 +12,7 @@ def stop_time(start_timer, stop_timer): torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) -def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, train_nodes): +def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_darts, train_nodes): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -104,6 +104,15 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr values=torch.cuda.LongTensor(p._nnz()).fill_(0), size=(mb_count, node_count)) sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + + neighbor_sizes = torch.sparse_coo_tensor(indices=p._indices(), + values=torch.cuda.LongTensor(p._nnz()).fill_(1), + size=(mb_count, node_count)) + neighbor_sizes = torch.sparse.sum(neighbor_sizes, dim=1) + + psum_neighbor_sizes = torch.cumsum(neighbor_sizes._values(), dim=0).roll(1) + psum_neighbor_sizes[0] = 0 + torch.cuda.nvtx.range_pop() print(f"pre-loop: {stop_time(start_timer, stop_timer)}") @@ -123,19 +132,24 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, tr start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") - dart_throws_values = torch.cuda.FloatTensor(p._nnz()).uniform_() + dartx_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() + darty_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - dart_throws_values = dart_throws_values * torch.gather(max_prob, 0, p._indices()[0, :]) + compute_darts_gpu(dartx_values, darty_values, neighbor_sizes._values(), psum_neighbor_sizes, + max_prob, n_darts, mb_count) + torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") - dart_hits_mask = dart_throws_values < p._values() + dart_hits_mask = torch.cuda.LongTensor(p._nnz()).fill_(0) + throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_mask, n_darts, mb_count) + dart_hits_mask = dart_hits_mask.bool() torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index cbe6dca..5826170 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -337,14 +337,14 @@ def main(args): # do it once before timing current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, \ args.n_bulkmb, args.n_layers, \ - train_nid) + args.n_darts, train_nid) print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, \ args.n_bulkmb, args.n_layers, \ - train_nid) + args.n_darts, train_nid) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() @@ -466,6 +466,8 @@ def main(args): 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('--n-darts', default=10, type=int, + help='number of darts to throw per minibatch in LADIES sampling') args = parser.parse_args() print(args) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index c0f495a..82e438e 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -253,8 +253,84 @@ void downsample_gpu(const at::Tensor& q_values, 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 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]) { + atomicCAS((unsigned long long *)&h_values[dartx_val], 0L, 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") +} + 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"); 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"); } diff --git a/tests/test_ladies.py b/tests/test_ladies.py index e695832..213ae0c 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -161,6 +161,8 @@ def main(args): 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) From 54cbd8072a4ff4e4f41771428f5a506d842307b6 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 8 Mar 2022 16:33:23 -0800 Subject: [PATCH 030/178] Downsample from H matrix instead of Q --- cagnet/samplers/ladies.py | 41 ++++++++++++++++++++++++++------------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0f46910..b5fa103 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -149,38 +149,53 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_hits_mask = torch.cuda.LongTensor(p._nnz()).fill_(0) throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_mask, n_darts, mb_count) - dart_hits_mask = dart_hits_mask.bool() torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-probs") - p._values()[dart_hits_mask] = 0.0 - torch.cuda.nvtx.range_pop() - timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") - next_frontier_values = torch.logical_or(dart_hits_mask.int(), next_frontier._values()).long() - next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), - values=next_frontier_values, - size=(mb_count, node_count)) + next_frontier_values = torch.logical_or(dart_hits_mask, next_frontier._values()).long() + next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=next_frontier_values, + size=(mb_count, node_count)) torch.cuda.nvtx.range_pop() timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-count-samples") - sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() + sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-downsample") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - downsample_gpu(next_frontier._values(), next_frontier._indices()[0,:], overflow, next_frontier._nnz()) + downsample_gpu(dart_hits_mask, next_frontier_tmp._indices()[0,:], overflow, next_frontier_tmp._nnz()) torch.cuda.nvtx.range_pop() timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-add-to-frontier2") + next_frontier_values = torch.logical_or(dart_hits_mask, next_frontier._values()).long() + next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=next_frontier_values, + size=(mb_count, node_count)) + torch.cuda.nvtx.range_pop() + timing_dict["add-to-frontier2"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-count-samples2") + sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() + torch.cuda.nvtx.range_pop() + timing_dict["count_samples2"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-set-probs") + dart_hits_mask = dart_hits_mask.bool() + p._values()[dart_hits_mask] = 0.0 + torch.cuda.nvtx.range_pop() + timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() # nvtx-sampling-iter timing_dict["sampling_iters"].append(stop_time(sample_start_timer, sample_stop_timer)) From 0acaa80811e79289c322283f8efda7eff60ee24d Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 6 Apr 2022 18:23:31 -0700 Subject: [PATCH 031/178] Uses sort() for selection, works on cora nbulk=4 --- .gitignore | 1 + cagnet/samplers/ladies.py | 38 +++++++++++++++++- sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 57 ++++++++++++++++++++++----- 4 files changed, 85 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index 7ecaa69..867478c 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ data/ job_scripts/ outputs/ profs/ +sparse-extension/build.out sparse-extension/build/ sparse-extension/dist/ sparse-extension/sparse_coo_tensor_cpp.egg-info/ diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index b5fa103..f4210ab 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -147,8 +147,10 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") + dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) dart_hits_mask = torch.cuda.LongTensor(p._nnz()).fill_(0) - throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_mask, n_darts, mb_count) + throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_count, dart_hits_mask, \ + n_darts, mb_count) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) @@ -170,7 +172,39 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-downsample") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - downsample_gpu(dart_hits_mask, next_frontier_tmp._indices()[0,:], overflow, next_frontier_tmp._nnz()) + # downsample_gpu(dart_hits_mask, next_frontier_tmp._indices()[0,:], \ + # overflow, \ + # next_frontier_tmp._nnz()) + dart_hits_expvar = torch.cuda.FloatTensor(p._nnz()).uniform_() + dart_hits_expvar = torch.pow(dart_hits_expvar, torch.reciprocal(dart_hits_count)) + dart_miss_mask = dart_hits_count == 0 + dart_hits_expvar[dart_miss_mask] = 1 + max_expvar = torch.max(dart_hits_expvar) + dart_hits_expvar = dart_hits_expvar + max_expvar * next_frontier._indices()[0,:] + dart_hits_expvar_sorted, dart_hits_expvar_idxs = torch.sort(dart_hits_expvar) + + dart_hits_rows = torch.bincount(next_frontier._indices()[0,:]) + ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) + ps_dart_hits_rows[0] = 0 + + # print(f"before dart_hits_mask.sum: {dart_hits_mask.sum()}") + # print(f"before dart_hits_mask: {dart_hits_mask}") + # print(f"before dart_hits_count: {dart_hits_count}") + # print(f"before h_rows: {next_frontier._indices()[0,:]}") + # print(f"before h_rows.size: {next_frontier._indices()[0,:].size()}") + # print(f"before ps_dart_hits_rows: {ps_dart_hits_rows}") + # print(f"before overflow: {overflow}") + # print(f"before dart_hits_expvar: {dart_hits_expvar}") + # print(f"before dart_hits_expvar_sorted: {dart_hits_expvar_sorted}") + # print(f"before dart_hits_expvar_idxs: {dart_hits_expvar_idxs}") + # print(f"before nnz: {next_frontier_tmp._nnz()}") + downsample_gpu(dart_hits_mask, next_frontier._indices()[0,:], \ + ps_dart_hits_rows, \ + dart_hits_expvar_idxs, \ + overflow, \ + next_frontier_tmp._nnz()) + # print(f"after dart_hits_mask.sum: {dart_hits_mask.sum()}") + torch.cuda.nvtx.range_pop() timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index 68f5f7b..43d73c9 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', ext_modules=[cpp_extension.CUDAExtension('sparse_coo_tensor_cpp', ['sparse_coo_tensor.cu'], - extra_compile_args=["-lcusparse"])], + extra_compile_args=["-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/include/", "-lcusparse"])], cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 82e438e..0e96235 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -226,20 +226,50 @@ void spmm_gpu(const at::Tensor& A_rowindices, C.t_(); } -__global__ void DownSample(long *q_values, long *q_rows, int *overflow, int nnz) { +// __global__ void DownSample(long *q_values, long *q_rows, int *overflow, int nnz) { +// int id = blockIdx.x * blockDim.x + threadIdx.x; +// int stride = blockDim.x * gridDim.x; +// +// for (int i = id; i < nnz; i += stride) { +// long row = q_rows[i]; +// if (q_values[i] == 1 && ((int) atomicSub((unsigned int *)&overflow[row], 1)) > 0) { +// q_values[i] = 0; +// } +// } +// } + +// void downsample_gpu(const at::Tensor& q_values, +// const at::Tensor& q_rows, +// 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<<>>(q_values.data(), q_rows.data(), +// overflow.data(), nnz); +// CHECK_ERROR("downsampling error") +// } + +__global__ void DownSample(long *h_mask, long *h_rows, long *ps_h_rows, long *hev_indices, int *overflow, + int nnz) { + int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; for (int i = id; i < nnz; i += stride) { - long row = q_rows[i]; - if (q_values[i] == 1 && ((int) atomicSub((unsigned int *)&overflow[row], 1)) > 0) { - q_values[i] = 0; + if (i - ps_h_rows[h_rows[i]] < overflow[h_rows[i]]) { + h_mask[hev_indices[i]] = 0; } } } -void downsample_gpu(const at::Tensor& q_values, - const at::Tensor& q_rows, +void downsample_gpu(const at::Tensor& h_mask, + const at::Tensor& h_rows, + const at::Tensor& ps_h_rows, + const at::Tensor& hev_indices, const at::Tensor& overflow, int nnz) { @@ -248,8 +278,12 @@ void downsample_gpu(const at::Tensor& q_values, int BLOCK_COUNT = std::ceil(nnz / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - DownSample<<>>(q_values.data(), q_rows.data(), - overflow.data(), nnz); + DownSample<<>>(h_mask.data(), + h_rows.data(), + ps_h_rows.data(), + hev_indices.data(), + overflow.data(), + nnz); CHECK_ERROR("downsampling error") } @@ -292,7 +326,7 @@ void compute_darts_gpu(const at::Tensor& dartx_values, } __global__ void ThrowDarts(float *dartx_values, float *darty_values, float *p_values, - long *h_values, int n_darts, int mb_count) { + long *h_values, long *h_mask, int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -301,7 +335,8 @@ __global__ void ThrowDarts(float *dartx_values, float *darty_values, float *p_va for (int i = id; i < dart_count; i += stride) { long dartx_val = (long) dartx_values[i]; if (darty_values[i] < p_values[dartx_val]) { - atomicCAS((unsigned long long *)&h_values[dartx_val], 0L, 1L); + atomicCAS((unsigned long long *)&h_mask[dartx_val], 0L, 1L); + atomicAdd((unsigned long long *)&h_values[dartx_val], 1L); } } } @@ -310,6 +345,7 @@ void throw_darts_gpu(const at::Tensor& dartx_values, const at::Tensor& darty_values, const at::Tensor& p_values, const at::Tensor& h_values, + const at::Tensor& h_mask, int n_darts, int mb_count) { @@ -322,6 +358,7 @@ void throw_darts_gpu(const at::Tensor& dartx_values, darty_values.data(), p_values.data(), h_values.data(), + h_mask.data(), n_darts, mb_count); CHECK_ERROR("dart throwing error") From 140d8df02a11fc19473d49051d4b8882dd464500 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 6 Apr 2022 18:24:32 -0700 Subject: [PATCH 032/178] Cleanup --- cagnet/samplers/ladies.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index f4210ab..cbccbb8 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -172,9 +172,6 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-downsample") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - # downsample_gpu(dart_hits_mask, next_frontier_tmp._indices()[0,:], \ - # overflow, \ - # next_frontier_tmp._nnz()) dart_hits_expvar = torch.cuda.FloatTensor(p._nnz()).uniform_() dart_hits_expvar = torch.pow(dart_hits_expvar, torch.reciprocal(dart_hits_count)) dart_miss_mask = dart_hits_count == 0 @@ -187,23 +184,11 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) ps_dart_hits_rows[0] = 0 - # print(f"before dart_hits_mask.sum: {dart_hits_mask.sum()}") - # print(f"before dart_hits_mask: {dart_hits_mask}") - # print(f"before dart_hits_count: {dart_hits_count}") - # print(f"before h_rows: {next_frontier._indices()[0,:]}") - # print(f"before h_rows.size: {next_frontier._indices()[0,:].size()}") - # print(f"before ps_dart_hits_rows: {ps_dart_hits_rows}") - # print(f"before overflow: {overflow}") - # print(f"before dart_hits_expvar: {dart_hits_expvar}") - # print(f"before dart_hits_expvar_sorted: {dart_hits_expvar_sorted}") - # print(f"before dart_hits_expvar_idxs: {dart_hits_expvar_idxs}") - # print(f"before nnz: {next_frontier_tmp._nnz()}") downsample_gpu(dart_hits_mask, next_frontier._indices()[0,:], \ ps_dart_hits_rows, \ dart_hits_expvar_idxs, \ overflow, \ next_frontier_tmp._nnz()) - # print(f"after dart_hits_mask.sum: {dart_hits_mask.sum()}") torch.cuda.nvtx.range_pop() timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) From 3d25768cbe0d5b1b6219177fdb3a868628278e00 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 6 Apr 2022 18:30:11 -0700 Subject: [PATCH 033/178] Removed unnecssary dart_hits_mask --- cagnet/samplers/ladies.py | 12 +++++------- sparse-extension/sparse_coo_tensor.cu | 13 +++++-------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index cbccbb8..ec51fdd 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -148,15 +148,13 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) - dart_hits_mask = torch.cuda.LongTensor(p._nnz()).fill_(0) - throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_count, dart_hits_mask, \ - n_darts, mb_count) + throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_count, n_darts, mb_count) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") - next_frontier_values = torch.logical_or(dart_hits_mask, next_frontier._values()).long() + next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=next_frontier_values, size=(mb_count, node_count)) @@ -184,7 +182,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) ps_dart_hits_rows[0] = 0 - downsample_gpu(dart_hits_mask, next_frontier._indices()[0,:], \ + downsample_gpu(dart_hits_count, next_frontier._indices()[0,:], \ ps_dart_hits_rows, \ dart_hits_expvar_idxs, \ overflow, \ @@ -195,7 +193,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier2") - next_frontier_values = torch.logical_or(dart_hits_mask, next_frontier._values()).long() + next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=next_frontier_values, size=(mb_count, node_count)) @@ -210,7 +208,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") - dart_hits_mask = dart_hits_mask.bool() + dart_hits_mask = dart_hits_count > 0 p._values()[dart_hits_mask] = 0.0 torch.cuda.nvtx.range_pop() timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 0e96235..60ba330 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -253,7 +253,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, // CHECK_ERROR("downsampling error") // } -__global__ void DownSample(long *h_mask, long *h_rows, long *ps_h_rows, long *hev_indices, int *overflow, +__global__ void DownSample(long *h_counts, long *h_rows, long *ps_h_rows, long *hev_indices, int *overflow, int nnz) { int id = blockIdx.x * blockDim.x + threadIdx.x; @@ -261,12 +261,12 @@ __global__ void DownSample(long *h_mask, long *h_rows, long *ps_h_rows, long *he for (int i = id; i < nnz; i += stride) { if (i - ps_h_rows[h_rows[i]] < overflow[h_rows[i]]) { - h_mask[hev_indices[i]] = 0; + h_counts[hev_indices[i]] = 0; } } } -void downsample_gpu(const at::Tensor& h_mask, +void downsample_gpu(const at::Tensor& h_counts, const at::Tensor& h_rows, const at::Tensor& ps_h_rows, const at::Tensor& hev_indices, @@ -278,7 +278,7 @@ void downsample_gpu(const at::Tensor& h_mask, int BLOCK_COUNT = std::ceil(nnz / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - DownSample<<>>(h_mask.data(), + DownSample<<>>(h_counts.data(), h_rows.data(), ps_h_rows.data(), hev_indices.data(), @@ -326,7 +326,7 @@ void compute_darts_gpu(const at::Tensor& dartx_values, } __global__ void ThrowDarts(float *dartx_values, float *darty_values, float *p_values, - long *h_values, long *h_mask, int n_darts, int mb_count) { + long *h_values, int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -335,7 +335,6 @@ __global__ void ThrowDarts(float *dartx_values, float *darty_values, float *p_va for (int i = id; i < dart_count; i += stride) { long dartx_val = (long) dartx_values[i]; if (darty_values[i] < p_values[dartx_val]) { - atomicCAS((unsigned long long *)&h_mask[dartx_val], 0L, 1L); atomicAdd((unsigned long long *)&h_values[dartx_val], 1L); } } @@ -345,7 +344,6 @@ void throw_darts_gpu(const at::Tensor& dartx_values, const at::Tensor& darty_values, const at::Tensor& p_values, const at::Tensor& h_values, - const at::Tensor& h_mask, int n_darts, int mb_count) { @@ -358,7 +356,6 @@ void throw_darts_gpu(const at::Tensor& dartx_values, darty_values.data(), p_values.data(), h_values.data(), - h_mask.data(), n_darts, mb_count); CHECK_ERROR("dart throwing error") From a9b093f388f778a6619c1d003afc1034f45c11cc Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 11 Apr 2022 15:13:08 -0700 Subject: [PATCH 034/178] Store debugging commit --- cagnet/samplers/ladies.py | 13 +++++++++++++ sparse-extension/sparse_coo_tensor.cu | 2 ++ 2 files changed, 15 insertions(+) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index ec51fdd..788ee89 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -163,9 +163,12 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-count-samples") + print(f"next_frontier_tmp._values(): {next_frontier_tmp._values()}") sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) + print(f"sampled_count1: {sampled_count}") + print(f"sampled_count1[209]: {sampled_count[209]}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-downsample") @@ -182,11 +185,19 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) ps_dart_hits_rows[0] = 0 + print(f"before dart_hits_count: {dart_hits_count}") + print(f"before dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}") + print(f"before dart_hits_count[mb 209].nonzero().size: {(dart_hits_count[next_frontier._indices()[0,:] == 209]).nonzero().size()}") + print(f"overflow: {overflow}") + print(f"overflow[209]: {overflow[209]}") + print(f"overflow.sum: {overflow.sum()}") downsample_gpu(dart_hits_count, next_frontier._indices()[0,:], \ ps_dart_hits_rows, \ dart_hits_expvar_idxs, \ overflow, \ next_frontier_tmp._nnz()) + print(f"after dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}") + print(f"after dart_hits_count[mb 209].nonzero().size: {(dart_hits_count[next_frontier._indices()[0,:] == 209]).nonzero().size()}") torch.cuda.nvtx.range_pop() timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) @@ -205,6 +216,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples2"].append(stop_time(start_timer, stop_timer)) + print(f"sampled_count2: {sampled_count}") + print(f"sampled_count2[209]: {sampled_count[209]}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 60ba330..4f1e5eb 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -262,6 +262,8 @@ __global__ void DownSample(long *h_counts, long *h_rows, long *ps_h_rows, long * for (int i = id; i < nnz; i += stride) { if (i - ps_h_rows[h_rows[i]] < overflow[h_rows[i]]) { h_counts[hev_indices[i]] = 0; + } else if (h_rows[i] == 209) { + printf("i: %d hev_indices[i]: %ld h_counts[hev_indices[i]]: %ld\n", i, hev_indices[i], h_counts[hev_indices[i]]); } } } From 41a31d9903314e00afbee77b8706db6925b7be6f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 14 Apr 2022 16:31:01 -0700 Subject: [PATCH 035/178] Fixed max_expvr bug --- cagnet/samplers/ladies.py | 32 +++++++++++++++++++++++---- sparse-extension/sparse_coo_tensor.cu | 10 ++++----- 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 788ee89..c68cea4 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -163,7 +163,8 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-count-samples") - print(f"next_frontier_tmp._values(): {next_frontier_tmp._values()}") + # print(f"next_frontier_tmp._values()[mb 209]: {next_frontier_tmp._values()[next_frontier._indices()[0,:] == 209].size()}") + # print(f"next_frontier_tmp._values()[mb 209]: {next_frontier_tmp._values()[next_frontier._indices()[0,:] == 209].nonzero().size()}") sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) @@ -174,11 +175,20 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ torch.cuda.nvtx.range_push("nvtx-downsample") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() dart_hits_expvar = torch.cuda.FloatTensor(p._nnz()).uniform_() + print(f"dart_hits_expvar == 0.size: {(dart_hits_expvar == 0).nonzero().size()}") + print(f"dart_hits_expvar == 1.size: {(dart_hits_expvar == 1).nonzero().size()}") + print(f"before1 mod dart_hits_expvar[mb 209][580]: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209])[580]}") dart_hits_expvar = torch.pow(dart_hits_expvar, torch.reciprocal(dart_hits_count)) dart_miss_mask = dart_hits_count == 0 - dart_hits_expvar[dart_miss_mask] = 1 + dart_hits_expvar[dart_miss_mask] = 1.1 max_expvar = torch.max(dart_hits_expvar) + print(f"before2 mod dart_hits_expvar[mb 209][580]: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209])[580]}") dart_hits_expvar = dart_hits_expvar + max_expvar * next_frontier._indices()[0,:] + print(f"after mod dart_hits_expvar[mb 209][580]: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209])[580]}") + print(f"dart_hits_expvar[mb 209]: {dart_hits_expvar[next_frontier._indices()[0,:] == 209]}") + print(f"dart_hits_expvar[mb 209].size: {dart_hits_expvar[next_frontier._indices()[0,:] == 209].size()}") + print(f"dart_hits_expvar[mb 209].nonzero().size: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209] != 210).nonzero().size()}") + print(f"dart_hits_expvar[mb 209].nonzero(): {(dart_hits_expvar[next_frontier._indices()[0,:] == 209] != 210).nonzero()}") dart_hits_expvar_sorted, dart_hits_expvar_idxs = torch.sort(dart_hits_expvar) dart_hits_rows = torch.bincount(next_frontier._indices()[0,:]) @@ -187,16 +197,30 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ print(f"before dart_hits_count: {dart_hits_count}") print(f"before dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}") - print(f"before dart_hits_count[mb 209].nonzero().size: {(dart_hits_count[next_frontier._indices()[0,:] == 209]).nonzero().size()}") + mb209_dart_hits = dart_hits_count[next_frontier._indices()[0,:] == 209] + print(f"before dart_hits_count[mb 209]: {mb209_dart_hits}") + print(f"before dart_hits_count[mb 209][580]: {mb209_dart_hits[580]}") + print(f"before dart_hits_count[mb 209].nonzero().size(): {mb209_dart_hits.nonzero().size()}") + print(f"before dart_hits_count[mb 209].nonzero(): {mb209_dart_hits.nonzero()}") + # print(f"before dart_hits_count[mb 209].sum(): {mb209_dart_hits.sum()}") + # print(f"before dart_hits_count[mb 209].nonzero().size: {mb209_dart_hits.nonzero().size()}") + # print(f"before dart_hits_expvar_idxs[mb 209]: {dart_hits_expvar_idxs[next_frontier._indices()[0,:] == 209]}") + # print(f"before dart_hits_expvar_idxs[mb 209].idx: {(dart_hits_expvar_idxs[next_frontier._indices()[0,:] == 209] == 389878).nonzero()}") + # print(f"before dart_hits_expvar_idxs[mb 209].idx[613]: {(dart_hits_expvar_idxs[next_frontier._indices()[0,:] == 209])[613]}") + # print(f"before dart_hits_expvar_sorted[mb 209].idx[612]: {(dart_hits_expvar_sorted[next_frontier._indices()[0,:] == 209])[612]}") + # print(f"before dart_hits_expvar_sorted[mb 209].idx[613]: {(dart_hits_expvar_sorted[next_frontier._indices()[0,:] == 209])[613]}") + # print(f"before dart_hits_expvar_sorted[mb 209].idx[614]: {(dart_hits_expvar_sorted[next_frontier._indices()[0,:] == 209])[614]}") print(f"overflow: {overflow}") print(f"overflow[209]: {overflow[209]}") print(f"overflow.sum: {overflow.sum()}") + print(f"dart_hits_rows[mb 209]: {dart_hits_rows[209]}") + print(f"ps_dart_hits_rows[mb 209]: {ps_dart_hits_rows[209]}", flush=True) downsample_gpu(dart_hits_count, next_frontier._indices()[0,:], \ ps_dart_hits_rows, \ dart_hits_expvar_idxs, \ overflow, \ next_frontier_tmp._nnz()) - print(f"after dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}") + print(f"after dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}", flush=True) print(f"after dart_hits_count[mb 209].nonzero().size: {(dart_hits_count[next_frontier._indices()[0,:] == 209]).nonzero().size()}") torch.cuda.nvtx.range_pop() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 4f1e5eb..89aaa35 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -256,14 +256,14 @@ void spmm_gpu(const at::Tensor& A_rowindices, __global__ void DownSample(long *h_counts, long *h_rows, long *ps_h_rows, long *hev_indices, int *overflow, int nnz) { - int id = blockIdx.x * blockDim.x + threadIdx.x; - int stride = blockDim.x * gridDim.x; + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; - for (int i = id; i < nnz; i += stride) { + 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; - } else if (h_rows[i] == 209) { - printf("i: %d hev_indices[i]: %ld h_counts[hev_indices[i]]: %ld\n", i, hev_indices[i], h_counts[hev_indices[i]]); + } else if (h_rows[i] == 209 && h_counts[hev_indices[i]] != 0) { + printf("i: %d hev_indices[i]: %ld ps_h_rows[h_rows[i]]: %ld overflow[h_rows[i]]: %ld h_counts[hev_indices[i]]: %ld\n", i, hev_indices[i], ps_h_rows[h_rows[i]], overflow[h_rows[i]], h_counts[hev_indices[i]]); } } } From 4f81dda953cc19a9dbf1f235adbf36f45f301d58 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 14 Apr 2022 16:35:59 -0700 Subject: [PATCH 036/178] Cleanup --- cagnet/samplers/ladies.py | 37 --------------------------- sparse-extension/sparse_coo_tensor.cu | 2 -- 2 files changed, 39 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index c68cea4..cd84740 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -163,65 +163,30 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-count-samples") - # print(f"next_frontier_tmp._values()[mb 209]: {next_frontier_tmp._values()[next_frontier._indices()[0,:] == 209].size()}") - # print(f"next_frontier_tmp._values()[mb 209]: {next_frontier_tmp._values()[next_frontier._indices()[0,:] == 209].nonzero().size()}") sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) - print(f"sampled_count1: {sampled_count}") - print(f"sampled_count1[209]: {sampled_count[209]}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-downsample") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() dart_hits_expvar = torch.cuda.FloatTensor(p._nnz()).uniform_() - print(f"dart_hits_expvar == 0.size: {(dart_hits_expvar == 0).nonzero().size()}") - print(f"dart_hits_expvar == 1.size: {(dart_hits_expvar == 1).nonzero().size()}") - print(f"before1 mod dart_hits_expvar[mb 209][580]: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209])[580]}") dart_hits_expvar = torch.pow(dart_hits_expvar, torch.reciprocal(dart_hits_count)) dart_miss_mask = dart_hits_count == 0 dart_hits_expvar[dart_miss_mask] = 1.1 max_expvar = torch.max(dart_hits_expvar) - print(f"before2 mod dart_hits_expvar[mb 209][580]: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209])[580]}") dart_hits_expvar = dart_hits_expvar + max_expvar * next_frontier._indices()[0,:] - print(f"after mod dart_hits_expvar[mb 209][580]: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209])[580]}") - print(f"dart_hits_expvar[mb 209]: {dart_hits_expvar[next_frontier._indices()[0,:] == 209]}") - print(f"dart_hits_expvar[mb 209].size: {dart_hits_expvar[next_frontier._indices()[0,:] == 209].size()}") - print(f"dart_hits_expvar[mb 209].nonzero().size: {(dart_hits_expvar[next_frontier._indices()[0,:] == 209] != 210).nonzero().size()}") - print(f"dart_hits_expvar[mb 209].nonzero(): {(dart_hits_expvar[next_frontier._indices()[0,:] == 209] != 210).nonzero()}") dart_hits_expvar_sorted, dart_hits_expvar_idxs = torch.sort(dart_hits_expvar) dart_hits_rows = torch.bincount(next_frontier._indices()[0,:]) ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) ps_dart_hits_rows[0] = 0 - print(f"before dart_hits_count: {dart_hits_count}") - print(f"before dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}") - mb209_dart_hits = dart_hits_count[next_frontier._indices()[0,:] == 209] - print(f"before dart_hits_count[mb 209]: {mb209_dart_hits}") - print(f"before dart_hits_count[mb 209][580]: {mb209_dart_hits[580]}") - print(f"before dart_hits_count[mb 209].nonzero().size(): {mb209_dart_hits.nonzero().size()}") - print(f"before dart_hits_count[mb 209].nonzero(): {mb209_dart_hits.nonzero()}") - # print(f"before dart_hits_count[mb 209].sum(): {mb209_dart_hits.sum()}") - # print(f"before dart_hits_count[mb 209].nonzero().size: {mb209_dart_hits.nonzero().size()}") - # print(f"before dart_hits_expvar_idxs[mb 209]: {dart_hits_expvar_idxs[next_frontier._indices()[0,:] == 209]}") - # print(f"before dart_hits_expvar_idxs[mb 209].idx: {(dart_hits_expvar_idxs[next_frontier._indices()[0,:] == 209] == 389878).nonzero()}") - # print(f"before dart_hits_expvar_idxs[mb 209].idx[613]: {(dart_hits_expvar_idxs[next_frontier._indices()[0,:] == 209])[613]}") - # print(f"before dart_hits_expvar_sorted[mb 209].idx[612]: {(dart_hits_expvar_sorted[next_frontier._indices()[0,:] == 209])[612]}") - # print(f"before dart_hits_expvar_sorted[mb 209].idx[613]: {(dart_hits_expvar_sorted[next_frontier._indices()[0,:] == 209])[613]}") - # print(f"before dart_hits_expvar_sorted[mb 209].idx[614]: {(dart_hits_expvar_sorted[next_frontier._indices()[0,:] == 209])[614]}") - print(f"overflow: {overflow}") - print(f"overflow[209]: {overflow[209]}") - print(f"overflow.sum: {overflow.sum()}") - print(f"dart_hits_rows[mb 209]: {dart_hits_rows[209]}") - print(f"ps_dart_hits_rows[mb 209]: {ps_dart_hits_rows[209]}", flush=True) downsample_gpu(dart_hits_count, next_frontier._indices()[0,:], \ ps_dart_hits_rows, \ dart_hits_expvar_idxs, \ overflow, \ next_frontier_tmp._nnz()) - print(f"after dart_hits_count.nonzero().size: {dart_hits_count.nonzero().size()}", flush=True) - print(f"after dart_hits_count[mb 209].nonzero().size: {(dart_hits_count[next_frontier._indices()[0,:] == 209]).nonzero().size()}") torch.cuda.nvtx.range_pop() timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) @@ -240,8 +205,6 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() torch.cuda.nvtx.range_pop() timing_dict["count_samples2"].append(stop_time(start_timer, stop_timer)) - print(f"sampled_count2: {sampled_count}") - print(f"sampled_count2[209]: {sampled_count[209]}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 89aaa35..f947fae 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -262,8 +262,6 @@ __global__ void DownSample(long *h_counts, long *h_rows, long *ps_h_rows, long * 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; - } else if (h_rows[i] == 209 && h_counts[hev_indices[i]] != 0) { - printf("i: %d hev_indices[i]: %ld ps_h_rows[h_rows[i]]: %ld overflow[h_rows[i]]: %ld h_counts[hev_indices[i]]: %ld\n", i, hev_indices[i], ps_h_rows[h_rows[i]], overflow[h_rows[i]], h_counts[hev_indices[i]]); } } } From 756ddf5fd8187303ef5ef7ab63fd0164f56829bd Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 15 Apr 2022 15:20:03 -0700 Subject: [PATCH 037/178] Split nvtx ranges --- cagnet/samplers/ladies.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index cd84740..7aadff9 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -168,20 +168,33 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-downsample") + torch.cuda.nvtx.range_push("nvtx-downsample-preproc") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() dart_hits_expvar = torch.cuda.FloatTensor(p._nnz()).uniform_() dart_hits_expvar = torch.pow(dart_hits_expvar, torch.reciprocal(dart_hits_count)) dart_miss_mask = dart_hits_count == 0 - dart_hits_expvar[dart_miss_mask] = 1.1 + dart_hits_expvar[dart_miss_mask] = 2 max_expvar = torch.max(dart_hits_expvar) dart_hits_expvar = dart_hits_expvar + max_expvar * next_frontier._indices()[0,:] + torch.cuda.nvtx.range_pop() + timing_dict["downsample-preproc"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-downsample-sort") dart_hits_expvar_sorted, dart_hits_expvar_idxs = torch.sort(dart_hits_expvar) + torch.cuda.nvtx.range_pop() + timing_dict["downsample-sort"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-downsample-psum") dart_hits_rows = torch.bincount(next_frontier._indices()[0,:]) ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) ps_dart_hits_rows[0] = 0 + torch.cuda.nvtx.range_pop() + timing_dict["downsample-psum"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-downsample-kernel") downsample_gpu(dart_hits_count, next_frontier._indices()[0,:], \ ps_dart_hits_rows, \ dart_hits_expvar_idxs, \ @@ -189,7 +202,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ next_frontier_tmp._nnz()) torch.cuda.nvtx.range_pop() - timing_dict["downsample"].append(stop_time(start_timer, stop_timer)) + timing_dict["downsample-kernel"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier2") From 86e19ad9985f769590ee2d21208774085561a770 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 21 Apr 2022 00:28:54 -0700 Subject: [PATCH 038/178] Implemented dart throwing selection --- cagnet/samplers/ladies.py | 98 ++++++++++++++----------- sparse-extension/sparse_coo_tensor.cu | 101 ++++++++++++++++++++++++++ 2 files changed, 155 insertions(+), 44 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 7aadff9..fe01abb 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -2,7 +2,8 @@ import torch_sparse from collections import defaultdict -from sparse_coo_tensor_cpp import downsample_gpu, compute_darts_gpu, throw_darts_gpu +from sparse_coo_tensor_cpp import downsample_gpu, compute_darts_gpu, throw_darts_gpu, compute_darts_select_gpu, \ + throw_darts_select_gpu def start_time(timer): timer.record() @@ -168,61 +169,70 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-downsample-preproc") + torch.cuda.nvtx.range_push("nvtx-dart-selection") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - dart_hits_expvar = torch.cuda.FloatTensor(p._nnz()).uniform_() - dart_hits_expvar = torch.pow(dart_hits_expvar, torch.reciprocal(dart_hits_count)) - dart_miss_mask = dart_hits_count == 0 - dart_hits_expvar[dart_miss_mask] = 2 - max_expvar = torch.max(dart_hits_expvar) - dart_hits_expvar = dart_hits_expvar + max_expvar * next_frontier._indices()[0,:] - torch.cuda.nvtx.range_pop() - timing_dict["downsample-preproc"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-downsample-sort") - dart_hits_expvar_sorted, dart_hits_expvar_idxs = torch.sort(dart_hits_expvar) - torch.cuda.nvtx.range_pop() - timing_dict["downsample-sort"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-downsample-psum") - dart_hits_rows = torch.bincount(next_frontier._indices()[0,:]) - ps_dart_hits_rows = torch.cumsum(dart_hits_rows, dim=0).roll(1) - ps_dart_hits_rows[0] = 0 - torch.cuda.nvtx.range_pop() - timing_dict["downsample-psum"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-downsample-kernel") - downsample_gpu(dart_hits_count, next_frontier._indices()[0,:], \ - ps_dart_hits_rows, \ - dart_hits_expvar_idxs, \ - overflow, \ - next_frontier_tmp._nnz()) + while (overflow > 0).any(): + print(f"overflow: {overflow}") + ps_overflow = torch.cumsum(overflow, dim=0) + total_overflow = ps_overflow[-1].item() + # ps_overflow = ps_overflow.roll(1) + # ps_overflow[0] = 0 + + dart_hits_inv = dart_hits_count.reciprocal() + dart_hits_inv[dart_hits_inv == float("inf")] = 0 + # print(f"dart_hits_count: {dart_hits_count}") + # print(f"dart_hits_inv: {dart_hits_inv}") + dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=dart_hits_inv, + size=(mb_count, node_count)) + + dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() + ps_dart_hits_inv_sum = torch.cumsum(dart_hits_inv_sum, dim=0).roll(1) + ps_dart_hits_inv_sum[0] = 0 + + dart_select = torch.cuda.FloatTensor(total_overflow).uniform_() + + # Compute darts for selection + # print(f"before dart_select: {dart_select}") + # print(f"before dart_hits_inv_sum: {dart_hits_inv_sum}") + # print(f"before ps_dart_hits_inv_sum: {ps_dart_hits_inv_sum}", flush=True) + compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, + mb_count, total_overflow) + # print(f"after dart_select: {dart_select}", flush=True) + + # Throw selection darts + ps_dart_hits_inv = torch.cumsum(dart_hits_inv, dim=0) + # ps_dart_hits_inv = ps_dart_hits_inv.roll(1) + # ps_dart_hits_inv[0] = 0 + # print(f"before ps_dart_hits_inv: {ps_dart_hits_inv}") + # print(f"before ps_dart_hits_inv.size: {ps_dart_hits_inv.size()}") + # print(f"before dart_hits_count: {dart_hits_count}", flush=True) + throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count, total_overflow, + dart_hits_inv_mtx._nnz()) + # print(f"after dart_hits_count: {dart_hits_count}", flush=True) + + next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() + next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=next_frontier_values, + size=(mb_count, node_count)) + + sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() torch.cuda.nvtx.range_pop() - timing_dict["downsample-kernel"].append(stop_time(start_timer, stop_timer)) + timing_dict["dart-selection"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-add-to-frontier2") next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), - values=next_frontier_values, - size=(mb_count, node_count)) - torch.cuda.nvtx.range_pop() - timing_dict["add-to-frontier2"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-count-samples2") - sampled_count = torch.sparse.sum(next_frontier, dim=1)._values() - torch.cuda.nvtx.range_pop() - timing_dict["count_samples2"].append(stop_time(start_timer, stop_timer)) + values=next_frontier_values, + size=(mb_count, node_count)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") dart_hits_mask = dart_hits_count > 0 p._values()[dart_hits_mask] = 0.0 + # TODO: Set probability values to 0 for minibatches sufficiently sampled torch.cuda.nvtx.range_pop() timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index f947fae..3887ee1 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -325,6 +325,57 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } +__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 { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +__global__ void ComputeDartsSelect(float *dart_select, float *dart_hits_inv_sum, float *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); + + 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) { @@ -361,10 +412,60 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } +__device__ long binary_searchf(float *arr, float val, long imin, long imax) { + + long ans = -1; + while (imax >= imin) { + long imid = (imin + imax) / 2; + + if (arr[imid] <= val) { + imin = imid + 1; + } else { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} + +__global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, long *dart_hits_count, + long total_overflow, long nnz) { + + long id = blockIdx.x * blockDim.x + threadIdx.x; + long stride = blockDim.x * gridDim.x; + + for (long i = id; i < total_overflow; i += stride) { + long vtx = binary_searchf(ps_dart_hits_inv, dart_select[i], 0, nnz); + atomicAnd((unsigned long long *)&dart_hits_count[vtx], 0L); + } +} + +void throw_darts_select_gpu(const at::Tensor& dart_select, + const at::Tensor& ps_dart_hits_inv, + const at::Tensor& dart_hits_count, + long total_overflow, + long nnz) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + ThrowDartsSelect<<>>(dart_select.data(), + ps_dart_hits_inv.data(), + dart_hits_count.data(), + total_overflow, + nnz); + CHECK_ERROR("selection dart throwing error") +} + 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"); 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"); } From 160cc1187f4d0cbaca97e2e97ff76eef2355d0c2 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 21 Apr 2022 00:30:04 -0700 Subject: [PATCH 039/178] Cleanup --- cagnet/samplers/ladies.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index fe01abb..a4c2ab9 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -176,13 +176,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ print(f"overflow: {overflow}") ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() - # ps_overflow = ps_overflow.roll(1) - # ps_overflow[0] = 0 dart_hits_inv = dart_hits_count.reciprocal() dart_hits_inv[dart_hits_inv == float("inf")] = 0 - # print(f"dart_hits_count: {dart_hits_count}") - # print(f"dart_hits_inv: {dart_hits_inv}") dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=dart_hits_inv, size=(mb_count, node_count)) @@ -194,23 +190,13 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ dart_select = torch.cuda.FloatTensor(total_overflow).uniform_() # Compute darts for selection - # print(f"before dart_select: {dart_select}") - # print(f"before dart_hits_inv_sum: {dart_hits_inv_sum}") - # print(f"before ps_dart_hits_inv_sum: {ps_dart_hits_inv_sum}", flush=True) compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) - # print(f"after dart_select: {dart_select}", flush=True) # Throw selection darts ps_dart_hits_inv = torch.cumsum(dart_hits_inv, dim=0) - # ps_dart_hits_inv = ps_dart_hits_inv.roll(1) - # ps_dart_hits_inv[0] = 0 - # print(f"before ps_dart_hits_inv: {ps_dart_hits_inv}") - # print(f"before ps_dart_hits_inv.size: {ps_dart_hits_inv.size()}") - # print(f"before dart_hits_count: {dart_hits_count}", flush=True) throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count, total_overflow, dart_hits_inv_mtx._nnz()) - # print(f"after dart_hits_count: {dart_hits_count}", flush=True) next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), From 3867512b504f2be5bb3b8cec24bf60c08b08c362 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 22 Apr 2022 00:15:34 -0700 Subject: [PATCH 040/178] Replaced some torch sparse sum instances --- cagnet/samplers/ladies.py | 71 +++++++++++++++++++++++++++++++++------ 1 file changed, 60 insertions(+), 11 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index a4c2ab9..dd0b29f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -20,6 +20,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ 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) + total_start_timer = torch.cuda.Event(enable_timing=True) total_stop_timer = torch.cuda.Event(enable_timing=True) @@ -169,45 +172,88 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-dart-selection") + torch.cuda.nvtx.range_push("nvtx-select-preproc") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + overflowed_minibatches = (overflow > 0).any() + selection_iter_count = 0 + timing_dict["select-preproc"].append(stop_time(start_timer, stop_timer)) + + while overflowed_minibatches: + start_time(select_start_timer) + torch.cuda.nvtx.range_push("nvtx-dart-selection") - while (overflow > 0).any(): - print(f"overflow: {overflow}") + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-psoverflow") + selection_iter_count += 1 ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() + timing_dict["select-psoverflow"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-reciprocal") dart_hits_inv = dart_hits_count.reciprocal() dart_hits_inv[dart_hits_inv == float("inf")] = 0 + timing_dict["select-reciprocal"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-instmtx") dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=dart_hits_inv, size=(mb_count, node_count)) + timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-invsum") + # dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() + dart_hits_inv_sum = torch.cuda.FloatTensor(mb_count).fill_(0) + dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) + timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) - dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-psinv") 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) + torch.cuda.nvtx.range_push("nvtx-select-computedarts") dart_select = torch.cuda.FloatTensor(total_overflow).uniform_() - # Compute darts for selection compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) + timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-throwdarts") # 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, - dart_hits_inv_mtx._nnz()) + next_frontier._nnz()) + timing_dict["select-throwdarts"].append(stop_time(start_timer, stop_timer)) - next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") + 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=(mb_count, node_count)) + timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-samplecount") + # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values()) + timing_dict["select-samplecount"].append(stop_time(start_timer, stop_timer)) - sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-overflow") 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)) - torch.cuda.nvtx.range_pop() - timing_dict["dart-selection"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() + timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), @@ -218,7 +264,10 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ torch.cuda.nvtx.range_push("nvtx-set-probs") dart_hits_mask = dart_hits_count > 0 p._values()[dart_hits_mask] = 0.0 - # TODO: Set probability values to 0 for minibatches sufficiently sampled + + filled_minibatches = sampled_count == frontier_size + filled_minibatches_mask = torch.gather(filled_minibatches, 0, p._indices()[0,:]) + p._values()[filled_minibatches_mask] = 0 torch.cuda.nvtx.range_pop() timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) From d56de442e6849c32b06a0cf9af57efd514076a08 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 25 Apr 2022 17:33:55 -0700 Subject: [PATCH 041/178] Replaced other sparse sum instances --- cagnet/samplers/ladies.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index dd0b29f..a9396a0 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -23,6 +23,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ 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) + total_start_timer = torch.cuda.Event(enable_timing=True) total_stop_timer = torch.cuda.Event(enable_timing=True) @@ -121,6 +124,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ print(f"pre-loop: {stop_time(start_timer, stop_timer)}") iter_count = 0 + selection_iter_count = 0 torch.cuda.nvtx.range_push("nvtx-sampling") while (sampled_count < frontier_size).any(): iter_count += 1 @@ -158,7 +162,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") - next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() + 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=(mb_count, node_count)) @@ -167,7 +171,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-count-samples") - sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + sampled_count = sampled_count.fill_(0) + sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0, :], next_frontier_values) torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) @@ -175,12 +181,14 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ torch.cuda.nvtx.range_push("nvtx-select-preproc") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() overflowed_minibatches = (overflow > 0).any() - selection_iter_count = 0 timing_dict["select-preproc"].append(stop_time(start_timer, stop_timer)) + start_time(select_start_timer) + torch.cuda.nvtx.range_push("nvtx-dart-selection") + while overflowed_minibatches: - start_time(select_start_timer) - torch.cuda.nvtx.range_push("nvtx-dart-selection") + start_time(select_iter_start_timer) + torch.cuda.nvtx.range_push("nvtx-selection-iter") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-psoverflow") @@ -252,8 +260,10 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ overflowed_minibatches = (overflow > 0).any() timing_dict["select-overflow"].append(stop_time(start_timer, stop_timer)) - torch.cuda.nvtx.range_pop() - timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) + timing_dict["select-iter"].append(stop_time(select_iter_start_timer, select_iter_stop_timer)) + + torch.cuda.nvtx.range_pop() # nvtx-dart-selection + timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), @@ -339,4 +349,5 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ for k, v in timing_dict.items(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") print(f"iter_count: {iter_count}") + print(f"selection_iter_count: {selection_iter_count}") return batches, next_frontier, adj_matrices From e22ddcdfbb88289eb79daa3000e26e01fb2ef03f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 26 Apr 2022 12:41:00 -0700 Subject: [PATCH 042/178] Replaced 2D rejection sampling with 1D sampling --- cagnet/samplers/ladies.py | 24 +++--- sparse-extension/sparse_coo_tensor.cu | 112 ++++++++++++++++++++------ 2 files changed, 101 insertions(+), 35 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index a9396a0..0574211 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -3,7 +3,7 @@ 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 + throw_darts_select_gpu, compute_darts1d_gpu, throw_darts1d_gpu def start_time(timer): timer.record() @@ -126,37 +126,39 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ iter_count = 0 selection_iter_count = 0 torch.cuda.nvtx.range_push("nvtx-sampling") + p_rowsum = torch.cuda.FloatTensor(mb_count).fill_(0) while (sampled_count < frontier_size).any(): iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-comp-maxprob") - p_torch_sparse = torch_sparse.tensor.SparseTensor.from_torch_sparse_coo_tensor(p) - max_prob = torch_sparse.reduce.reduction(p_torch_sparse, dim=1, reduce="max") + torch.cuda.nvtx.range_push("nvtx-prob-rowsum") + 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 + ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) + ps_p_values[0] = 0 torch.cuda.nvtx.range_pop() - timing_dict["max-prob"].append(stop_time(start_timer, stop_timer)) + timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") - dartx_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() - darty_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() + dart_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - compute_darts_gpu(dartx_values, darty_values, neighbor_sizes._values(), psum_neighbor_sizes, - max_prob, n_darts, mb_count) - + compute_darts1d_gpu(dart_values, p_rowsum, ps_p_rowsum, n_darts, mb_count) torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) - throw_darts_gpu(dartx_values, darty_values, p._values(), dart_hits_count, n_darts, mb_count) + throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 3887ee1..39a3bd6 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -31,6 +31,39 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} +__device__ long binary_searchf(float *arr, float val, long imin, long imax) { + + long ans = -1; + while (imax >= imin) { + long imid = (imin + imax) / 2; + + 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 { + ans = imid; + imax = imid - 1; + } + } + + return ans; +} at::Tensor expand_values_if_needed(const at::Tensor& values) { // expand @@ -325,21 +358,37 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } -__device__ long binary_searchl(long *arr, long val, long imin, long imax) { - - long ans = -1; - while (imax >= imin) { - long imid = (imin + imax) / 2; +__global__ void ComputeDarts1D(float *dart_values, float *p_rowsum, float *ps_p_rowsum, + int n_darts, int mb_count) { - if (arr[imid] <= val) { - imin = imid + 1; - } else { - ans = imid; - imax = imid - 1; - } + 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) { + int row = i / n_darts; + dart_values[i] *= p_rowsum[row]; + dart_values[i] += ps_p_rowsum[row]; } +} - return ans; +void compute_darts1d_gpu(const at::Tensor& dart_values, + const at::Tensor& p_rowsum, + const at::Tensor& ps_p_rowsum, + 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); + + ComputeDarts1D<<>>(dart_values.data(), + p_rowsum.data(), + ps_p_rowsum.data(), + n_darts, + mb_count); + CHECK_ERROR("dart1d computation error") } __global__ void ComputeDartsSelect(float *dart_select, float *dart_hits_inv_sum, float *ps_dart_hits_inv_sum, @@ -412,21 +461,34 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__device__ long binary_searchf(float *arr, float val, long imin, long imax) { - - long ans = -1; - while (imax >= imin) { - long imid = (imin + imax) / 2; +__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, long *h_values, int dart_count, int nnz) { - if (arr[imid] <= val) { - imin = imid + 1; - } else { - ans = imid; - imax = imid - 1; - } + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < dart_count; i += stride) { + long vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz); + atomicAdd((unsigned long long *)&h_values[vtx], 1L); } +} - return ans; +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(), + h_values.data(), + dart_count, + nnz); + CHECK_ERROR("dart throwing error") } __global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, long *dart_hits_count, @@ -468,4 +530,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From 401b35557b1d53d392ee0421a20a7e23f0351903 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 26 Apr 2022 17:27:29 -0700 Subject: [PATCH 043/178] Moved batch creation to main file --- cagnet/samplers/ladies.py | 23 ++--------------------- examples/gcn_15d.py | 29 +++++++++++++++++++++-------- 2 files changed, 23 insertions(+), 29 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0574211..40b2ae6 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -13,7 +13,7 @@ def stop_time(start_timer, stop_timer): torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) -def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_darts, train_nodes): +def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_layers, n_darts): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -31,28 +31,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ timing_dict = defaultdict(list) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-instantiations") - # frontiers = torch.cuda.IntTensor(n_layers, batch_size, mb_count) - batches = torch.cuda.IntTensor(mb_count, batch_size) # initially the minibatch, note row-major - current_frontier = torch.cuda.IntTensor(mb_count, batch_size + frontier_size) - # frontiers = torch.cuda.IntTensor(n_layers - 1, batch_size + frontier_size, mb_count) adj_matrices = [[None] * n_layers for x in range(mb_count)] # adj_matrices[i][j] -- mb i layer j + current_frontier = torch.cuda.IntTensor(mb_count, batch_size + frontier_size) node_count = adj_matrix.size(0) - torch.cuda.nvtx.range_pop() - print(f"instantiations: {stop_time(start_timer, stop_timer)}") - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - - torch.manual_seed(0) - vertex_perm = torch.randperm(train_nodes.size(0)) - # Generate minibatch vertices - for i in range(mb_count): - idx = vertex_perm[(i * batch_size):((i + 1) * batch_size)] - batches[i,:] = train_nodes[idx] - torch.cuda.nvtx.range_pop() - print(f"get-minibatch-vtxs: {stop_time(start_timer, stop_timer)}") start_time(total_start_timer) for i in range(n_layers): diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 5826170..f57d386 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -317,12 +317,11 @@ def main(args): ampbyp[i] = ampbyp[i].t().coalesce().to(device) print("coalescing", flush=True) - g_loc = g_loc.coalesce() + g_loc = g_loc.coalesce().t() print("normalizing", flush=True) g_loc = row_normalize(g_loc) print("done normalizing", flush=True) - n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] @@ -334,17 +333,31 @@ def main(args): train_nid = data.train_mask.nonzero().squeeze() test_nid = data.test_mask.nonzero().squeeze() + batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major + node_count = adj_matrix.size(0) + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") + torch.manual_seed(0) + 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() + + print(f"rank: {rank} g_loc: {g_loc}") # do it once before timing - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, \ - args.n_bulkmb, args.n_layers, \ - args.n_darts, train_nid) + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts) print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, args.batch_size, args.samp_num, \ - args.n_bulkmb, args.n_layers, \ - args.n_darts, train_nid) + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() From e482b17f64653a9547dc2ec6bfa5285a2480c9da Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 26 Apr 2022 19:35:18 -0700 Subject: [PATCH 044/178] Added all-gather for row/col selection temporarily --- cagnet/samplers/ladies.py | 77 ++++++++++++++++++++++----------------- examples/gcn_15d.py | 28 ++++++++++++-- 2 files changed, 68 insertions(+), 37 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 40b2ae6..04015a0 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -1,4 +1,5 @@ import torch +import torch.distributed as dist import torch_sparse from collections import defaultdict @@ -13,7 +14,9 @@ def stop_time(start_timer, stop_timer): torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) -def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_layers, n_darts): +def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ + rank, size, group): + start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -31,40 +34,28 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_l timing_dict = defaultdict(list) - adj_matrices = [[None] * n_layers for x in range(mb_count)] # adj_matrices[i][j] -- mb i layer j - current_frontier = torch.cuda.IntTensor(mb_count, batch_size + frontier_size) + adj_matrices = [[None] * n_layers for x in range(mb_count_total)] # adj_matrices[i][j] -- mb i layer j + 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) start_time(total_start_timer) for i in range(n_layers): if i == 0: - nnz = batches[0, :].size(0) + nnz = batch_size else: nnz = current_frontier[0, :].size(0) # A * Q^i # indices would change based on mb_count - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-gen-sparse-frontier") - if i == 0: - frontier_indices_rows = torch.arange(mb_count).cuda() - frontier_indices_rows = frontier_indices_rows.repeat_interleave(nnz) - frontier_indices_cols = batches.view(-1) - frontier_indices = torch.stack((frontier_indices_rows, frontier_indices_cols)) - else: - frontier_indices_rows = torch.arange(mb_count).cuda() - frontier_indices_rows = frontier_indices_rows.repeat_interleave(nnz) - frontier_indices_cols = current_frontier.view(-1) - frontier_indices = torch.stack((frontier_indices_rows, frontier_indices_cols)) - frontier_values = torch.cuda.FloatTensor(nnz * mb_count).fill_(1.0) - adj_mat_squared = torch.pow(adj_matrix._values(), 2) - torch.cuda.nvtx.range_pop() - print(f"gen-sparse-frontier: {stop_time(start_timer, stop_timer)}") - + + # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = torch_sparse.spspmm(frontier_indices.long(), frontier_values, - adj_matrix._indices().long(), adj_mat_squared, + p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), + adj_matrix._indices().long(), adj_matrix._values(), mb_count, node_count, node_count, coalesced=True) torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -274,9 +265,29 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_l start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-construct-nextf") - next_frontier = torch.masked_select(next_frontier._indices()[1,:], \ + next_frontier_select = torch.masked_select(next_frontier._indices()[1,:], \ next_frontier._values().bool()).view(mb_count, frontier_size) - next_frontier = torch.cat((next_frontier, batches), dim=1) + batches_select = torch.masked_select(batches._indices()[1,:], \ + batches._values().bool()).view(mb_count, batch_size) + next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) + + # next_frontier_select_recv = [] + # for j in range(size): + # next_frontier_select_recv.append(torch.cuda.LongTensor(next_frontier_select.size())) + # dist.all_gather(next_frontier_select_recv, next_frontier_select, group) + # next_frontier_select = torch.cat(next_frontier_select_recv, dim=0) + + # batches_select_recv = [] + # for j in range(size): + # batches_select_recv.append(torch.cuda.LongTensor(batches_select.size())) + # dist.all_gather(batches_select_recv, batches_select, group) + # batches_select = torch.cat(batches_select_recv, dim=0) + + print(f"rank: {rank} next_frontier_select.size: {next_frontier_select.size()}") + print(f"rank: {rank} batches_select.size: {batches_select.size()}") + print(f"rank: {rank} next_frontier_select: {next_frontier_select}") + print(f"rank: {rank} batches_select: {batches_select}") + torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") @@ -285,15 +296,15 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_l start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-mtxs") if i == 0: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches[j,:])) + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches_select[j,:])) else: row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), \ current_frontier[j, :])) row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) - col_select_mtx_indices = torch.stack((next_frontier[j], torch.arange(start=0, \ - end=next_frontier[j].size(0)).cuda())) - col_select_mtx_values = torch.cuda.FloatTensor(next_frontier[j].size(0)).fill_(1.0) + col_select_mtx_indices = torch.stack((next_frontier_select[j], torch.arange(start=0, \ + end=next_frontier_select[j].size(0)).cuda())) + col_select_mtx_values = torch.cuda.FloatTensor(next_frontier_select[j].size(0)).fill_(1.0) torch.cuda.nvtx.range_pop() timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) @@ -312,7 +323,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_l torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, col_select_mtx_indices.long(), col_select_mtx_values, - nnz, node_count, next_frontier[j].size(0), + nnz, node_count, next_frontier_select[j].size(0), coalesced=True) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -320,9 +331,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_l start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-sample") - current_frontier[j, :] = next_frontier[j, :] + current_frontier[j, :] = next_frontier_select[j, :] adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ - size=(nnz, next_frontier[j].size(0))) + size=(nnz, next_frontier_select[j].size(0))) adj_matrices[j][i] = adj_matrix_sample torch.cuda.nvtx.range_pop() timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) @@ -333,4 +344,4 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count, n_l print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") - return batches, next_frontier, adj_matrices + return batches_select, next_frontier_select, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index f57d386..1b5dc04 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -205,6 +205,11 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}", flush=True) return inputs_loc, adj_matrix_loc, am_pbyp +def oned_partition_mb(rank, size, batches, mb_count): + # TODO: Include replicaiton at some point + batch_partitions = torch.split(batches, math.ceil(float(mb_count / size)), dim=0) + return batch_partitions[rank] + def evaluate(model, graph, features, labels, nid, nid_count, ampbyp, group): model.eval() with torch.no_grad(): @@ -346,18 +351,33 @@ def main(args): batches[i,:] = train_nid[idx] torch.cuda.nvtx.range_pop() + batches_loc = 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) + torch.cuda.nvtx.range_pop() + print(f"rank: {rank} g_loc: {g_loc}") + print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches, args.batch_size, \ + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts) + args.n_layers, args.n_darts, \ + rank, size, col_groups[0]) print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches, args.batch_size, \ + current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts) + args.n_layers, args.n_darts, \ + rank, size, col_groups[0]) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() From de36970a3bd3c9087031c4dc17b22a9a3df0959a Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 28 Apr 2022 15:15:46 -0700 Subject: [PATCH 045/178] Added initial distributed spgemm implementation --- cagnet/samplers/ladies.py | 83 ++++++++++++++++++++++++++++++++++----- 1 file changed, 73 insertions(+), 10 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 04015a0..d84eae9 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -1,3 +1,4 @@ +import math import torch import torch.distributed as dist import torch_sparse @@ -14,6 +15,42 @@ def stop_time(start_timer, stop_timer): torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) +def dist_spgemm1D(mata, matb, rank, size, group): + chunk_size = math.ceil(float(mata.size(1) / size)) + matc = torch.sparse_coo_tensor(size=(mata.size(0), matb.size(1))).cuda() + + for i in range(size): + matb_recv_nnz = torch.cuda.IntTensor([matb._nnz()]) + dist.broadcast(matb_recv_nnz, src=i, group=group) + + if i == 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.FloatTensor(matb_recv_nnz.item()) + + dist.broadcast(matb_recv_indices, src=i, group=group) + dist.broadcast(matb_recv_values, src=i, group=group) + + chunk_col_start = i * chunk_size + chunk_col_stop = min((i + 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 -= chunk_col_start + mata_chunk_values = mata._values()[chunk_col_mask] + + 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 = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + matc += matc_chunk + + return matc._indices(), matc._values() + def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ rank, size, group): @@ -54,9 +91,35 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), - adj_matrix._indices().long(), adj_matrix._values(), - mb_count, node_count, node_count, coalesced=True) + print(f"rank: {rank} batches: {batches}") + print(f"rank: {rank} adj_matrix: {adj_matrix}") + # p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), + # adj_matrix._indices().long(), adj_matrix._values(), + # mb_count, node_count, node_count, coalesced=True) + p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) + # # all-gather so remaining computation is not distributed + # # p_num_indices[0, :] += rank * math.ceil(float(mb_count_total / size)) + # nnz_counts = [] + # for j in range(size): + # nnz_counts.append(torch.cuda.LongTensor(1)) + # dist.all_gather(nnz_counts, torch.cuda.LongTensor([p_num_values.size(0)]), group) + + # p_num_indices_recv = [] + # p_num_values_recv = [] + # for j in range(size): + # p_num_indices_recv.append(torch.cuda.LongTensor(2, nnz_counts[i].item())) + # p_num_values_recv.append(torch.cuda.FloatTensor(nnz_counts[i].item())) + + # p_num_indices_recv[rank] = p_num_indices + # p_num_values_recv[rank] = p_num_values + # for j in range(size): + # dist.broadcast(p_num_indices_recv[i], src=i, group=group) + # dist.broadcast(p_num_values_recv[i], src=i, group=group) + + # p_num_indices = torch.cat(p_num_indices_recv, dim=1) + # p_num_values = torch.cat(p_num_values_recv, dim=0) + # # mb_count = mb_count_total + # # node_count = node_count_total torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -66,7 +129,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ - size=(mb_count, node_count)) + size=(mb_count, node_count_total)) p_den = torch.sparse.sum(p_num, dim=1) for j in range(mb_count): @@ -81,12 +144,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-pre-loop") next_frontier = torch.sparse_coo_tensor(indices=p._indices(), values=torch.cuda.LongTensor(p._nnz()).fill_(0), - size=(mb_count, node_count)) + size=(mb_count, node_count_total)) sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) neighbor_sizes = torch.sparse_coo_tensor(indices=p._indices(), values=torch.cuda.LongTensor(p._nnz()).fill_(1), - size=(mb_count, node_count)) + size=(mb_count, node_count_total)) neighbor_sizes = torch.sparse.sum(neighbor_sizes, dim=1) psum_neighbor_sizes = torch.cumsum(neighbor_sizes._values(), dim=0).roll(1) @@ -139,7 +202,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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=(mb_count, node_count)) + size=(mb_count, node_count_total)) torch.cuda.nvtx.range_pop() timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) @@ -181,7 +244,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-select-instmtx") dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=dart_hits_inv, - size=(mb_count, node_count)) + size=(mb_count, node_count_total)) timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) @@ -218,7 +281,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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=(mb_count, node_count)) + size=(mb_count, node_count_total)) timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) @@ -242,7 +305,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() next_frontier = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=next_frontier_values, - size=(mb_count, node_count)) + size=(mb_count, node_count_total)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") From 7f04c574e8e2da26b7c521ef3382076ae2103a3a Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 30 Apr 2022 17:07:41 -0700 Subject: [PATCH 046/178] Initial test spgemm implementation --- cagnet/samplers/ladies.py | 4 +- test_nccl_p2p.py => tests/test_nccl_p2p.py | 0 tests/test_spgemm.py | 193 +++++++++++++++++++++ 3 files changed, 195 insertions(+), 2 deletions(-) rename test_nccl_p2p.py => tests/test_nccl_p2p.py (100%) create mode 100644 tests/test_spgemm.py diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index d84eae9..8c2b520 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -91,12 +91,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - print(f"rank: {rank} batches: {batches}") - print(f"rank: {rank} adj_matrix: {adj_matrix}") # p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), # adj_matrix._indices().long(), adj_matrix._values(), # mb_count, node_count, node_count, coalesced=True) p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) + print(f"rank: {rank} indices: {p_num_indices}") + print(f"rank: {rank} values: {p_num_values}") # # all-gather so remaining computation is not distributed # # p_num_indices[0, :] += rank * math.ceil(float(mb_count_total / size)) # nnz_counts = [] diff --git a/test_nccl_p2p.py b/tests/test_nccl_p2p.py similarity index 100% rename from test_nccl_p2p.py rename to tests/test_nccl_p2p.py diff --git a/tests/test_spgemm.py b/tests/test_spgemm.py new file mode 100644 index 0000000..043c6c2 --- /dev/null +++ b/tests/test_spgemm.py @@ -0,0 +1,193 @@ +import argparse +import examples.gcn_15d +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, size, group) + + 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}") + +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) From b6fdb28f25ad07fb5260fa87aae771cdad66516f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 30 Apr 2022 19:13:00 -0700 Subject: [PATCH 047/178] Working spgemm test script --- cagnet/samplers/ladies.py | 7 ++++--- tests/test_spgemm.py | 30 ++++++++++++++++++++++++++++-- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 8c2b520..ecf7d21 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -42,9 +42,10 @@ def dist_spgemm1D(mata, matb, rank, size, group): mata_chunk_indices -= chunk_col_start mata_chunk_values = mata._values()[chunk_col_mask] - 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_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 = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) matc += matc_chunk diff --git a/tests/test_spgemm.py b/tests/test_spgemm.py index 043c6c2..d93a50b 100644 --- a/tests/test_spgemm.py +++ b/tests/test_spgemm.py @@ -1,5 +1,6 @@ import argparse import examples.gcn_15d +import math import numpy as np import os import os.path as osp @@ -116,7 +117,7 @@ def main(args): # Run SpGEMM with only rank0, pre-partitioning if rank == 0: - ans_indices, ans_values = dist_spgemm1D(batches_loc, g_loc, 0, size, group) + ans_indices, ans_values = dist_spgemm1D(batches_loc, g_loc, 0, 1, row_groups[0]) dist.barrier(group) @@ -142,7 +143,7 @@ def main(args): 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 + # group = col_groups[0] # Assume replication factor == 1 dist.barrier(group) # Run SpGEMM with only rank0, pre-partitioning @@ -157,6 +158,31 @@ def main(args): 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", From 2df303b71b539974cecacbffeccb2286bb29040c Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 2 May 2022 12:13:20 -0700 Subject: [PATCH 048/178] Cleanup --- cagnet/samplers/ladies.py | 46 +++++++++++---------------------------- 1 file changed, 13 insertions(+), 33 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index ecf7d21..52cb18b 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -96,31 +96,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # adj_matrix._indices().long(), adj_matrix._values(), # mb_count, node_count, node_count, coalesced=True) p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) - print(f"rank: {rank} indices: {p_num_indices}") - print(f"rank: {rank} values: {p_num_values}") - # # all-gather so remaining computation is not distributed - # # p_num_indices[0, :] += rank * math.ceil(float(mb_count_total / size)) - # nnz_counts = [] - # for j in range(size): - # nnz_counts.append(torch.cuda.LongTensor(1)) - # dist.all_gather(nnz_counts, torch.cuda.LongTensor([p_num_values.size(0)]), group) - - # p_num_indices_recv = [] - # p_num_values_recv = [] - # for j in range(size): - # p_num_indices_recv.append(torch.cuda.LongTensor(2, nnz_counts[i].item())) - # p_num_values_recv.append(torch.cuda.FloatTensor(nnz_counts[i].item())) - - # p_num_indices_recv[rank] = p_num_indices - # p_num_values_recv[rank] = p_num_values - # for j in range(size): - # dist.broadcast(p_num_indices_recv[i], src=i, group=group) - # dist.broadcast(p_num_values_recv[i], src=i, group=group) - - # p_num_indices = torch.cat(p_num_indices_recv, dim=1) - # p_num_values = torch.cat(p_num_values_recv, dim=0) - # # mb_count = mb_count_total - # # node_count = node_count_total + torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -131,14 +107,18 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-compute-p") p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ size=(mb_count, node_count_total)) - p_den = torch.sparse.sum(p_num, dim=1) - - for j in range(mb_count): - if p_den._nnz() != mb_count: - print("ERROR nnz: {p_den._nnz()} mb_count: {mb_count}") - p_den_mb = p_den._values()[j].item() - p = torch.div(p_num, p_den_mb) - torch.cuda.nvtx.range_pop() + # p_den = torch.sparse.sum(p_num, dim=1) + p_den = torch.cuda.FloatTensor(mb_count).fill_(0) + p_den = p_den.scatter_add_(0, p_num._indices()[0, :], p_num._values()) + p_den = torch.reciprocal(p_den) + + # for j in range(mb_count): + # if p_den._nnz() != mb_count: + # print("ERROR nnz: {p_den._nnz()} mb_count: {mb_count}") + # p_den_mb = p_den._values()[j].item() + # p = torch.div(p_num, p_den_mb) + # torch.cuda.nvtx.range_pop() + p = torch.scatter_reduce(/ga print(f"compute-p: {stop_time(start_timer, stop_timer)}") start_time(start_timer) From 2e5eb9abc837e0d9563168f5c1ff2106db1de31d Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 2 May 2022 13:12:06 -0700 Subject: [PATCH 049/178] Normalize prob values in extension --- cagnet/samplers/ladies.py | 17 +++++------------ sparse-extension/sparse_coo_tensor.cu | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 52cb18b..781aaec 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -5,7 +5,8 @@ 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 + throw_darts_select_gpu, compute_darts1d_gpu, throw_darts1d_gpu, \ + normalize_gpu def start_time(timer): timer.record() @@ -105,20 +106,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p_num = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ + p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ size=(mb_count, node_count_total)) # p_den = torch.sparse.sum(p_num, dim=1) p_den = torch.cuda.FloatTensor(mb_count).fill_(0) - p_den = p_den.scatter_add_(0, p_num._indices()[0, :], p_num._values()) - p_den = torch.reciprocal(p_den) - - # for j in range(mb_count): - # if p_den._nnz() != mb_count: - # print("ERROR nnz: {p_den._nnz()} mb_count: {mb_count}") - # p_den_mb = p_den._values()[j].item() - # p = torch.div(p_num, p_den_mb) - # torch.cuda.nvtx.range_pop() - p = torch.scatter_reduce(/ga + p_den = p_den.scatter_add_(0, p._indices()[0, :], p._values()) + normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") start_time(start_timer) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 39a3bd6..31e2018 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -522,6 +522,26 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, CHECK_ERROR("selection dart throwing error") } +__global__ void Normalize(float *output, float *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); + CHECK_ERROR("normalize error") +} + 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"); @@ -532,4 +552,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From 8199dc4055d46db432e55ad3e9c6d6334f3c2b1e Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 2 May 2022 13:16:44 -0700 Subject: [PATCH 050/178] Removed unnecessary prefix sums --- cagnet/samplers/ladies.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 781aaec..75d6a70 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -93,22 +93,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - # p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), - # adj_matrix._indices().long(), adj_matrix._values(), - # mb_count, node_count, node_count, coalesced=True) p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) - torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") - print(f"p_num_indices.size(): {p_num_indices.size()}") - print(f"p_num_values.size(): {p_num_values.size()}") - start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ size=(mb_count, node_count_total)) - # p_den = torch.sparse.sum(p_num, dim=1) p_den = torch.cuda.FloatTensor(mb_count).fill_(0) p_den = p_den.scatter_add_(0, p._indices()[0, :], p._values()) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) @@ -121,14 +113,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota size=(mb_count, node_count_total)) sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) - neighbor_sizes = torch.sparse_coo_tensor(indices=p._indices(), - values=torch.cuda.LongTensor(p._nnz()).fill_(1), - size=(mb_count, node_count_total)) - neighbor_sizes = torch.sparse.sum(neighbor_sizes, dim=1) - - psum_neighbor_sizes = torch.cumsum(neighbor_sizes._values(), dim=0).roll(1) - psum_neighbor_sizes[0] = 0 - torch.cuda.nvtx.range_pop() print(f"pre-loop: {stop_time(start_timer, stop_timer)}") From 4476c15df2301c9367a8f2603f9c22222268563c Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 2 May 2022 18:38:41 -0700 Subject: [PATCH 051/178] Extracted underful_mb boolean --- cagnet/samplers/ladies.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 75d6a70..eee1ffd 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -120,7 +120,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota selection_iter_count = 0 torch.cuda.nvtx.range_push("nvtx-sampling") p_rowsum = torch.cuda.FloatTensor(mb_count).fill_(0) - while (sampled_count < frontier_size).any(): + + underfull_minibatches = (sampled_count < frontier_size).any() + while underfull_minibatches: iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -276,6 +278,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-compute-bool") + underfull_minibatches = (sampled_count < frontier_size).any() + torch.cuda.nvtx.range_pop() + timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() # nvtx-sampling-iter timing_dict["sampling_iters"].append(stop_time(sample_start_timer, sample_stop_timer)) From df7a73c4e463bb294e572fc083dbf00e21f9a0dc Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 3 May 2022 18:01:44 -0700 Subject: [PATCH 052/178] Fixed bug in distributed spgemm --- cagnet/samplers/ladies.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index eee1ffd..5c0c219 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -40,7 +40,7 @@ def dist_spgemm1D(mata, matb, rank, size, group): 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 -= chunk_col_start + mata_chunk_indices[1,:] -= chunk_col_start mata_chunk_values = mata._values()[chunk_col_mask] matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ @@ -93,14 +93,19 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") + print(f"batches: {batches}") + print(f"adj_matrices: {adj_matrix}") p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) + print(f"p_num_indices: {p_num_indices}") + print(f"p_num_values: {p_num_values}") torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, \ - size=(mb_count, node_count_total)) + p = torch.sparse_coo_tensor(indices=p_num_indices, + values=p_num_values, + size=(mb_count, node_count_total)) p_den = torch.cuda.FloatTensor(mb_count).fill_(0) p_den = p_den.scatter_add_(0, p._indices()[0, :], p._values()) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) @@ -121,7 +126,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-sampling") p_rowsum = torch.cuda.FloatTensor(mb_count).fill_(0) - underfull_minibatches = (sampled_count < frontier_size).any() + # underfull_minibatches = (sampled_count < frontier_size).any() + underfull_minibatches = True while underfull_minibatches: iter_count += 1 start_time(sample_start_timer) From 0e23548deb49795cbf7cf22445ae9fd904061059 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 4 May 2022 17:58:18 -0700 Subject: [PATCH 053/178] Bulk row/col selection for single mb --- cagnet/samplers/ladies.py | 144 +++++++++++++++----------- sparse-extension/sparse_coo_tensor.cu | 49 +++++++++ 2 files changed, 135 insertions(+), 58 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 5c0c219..e6d6ff1 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -6,7 +6,7 @@ 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_gpu, shift_rowselect_gpu, shift_colselect_gpu def start_time(timer): timer.record() @@ -306,18 +306,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches._values().bool()).view(mb_count, batch_size) next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) - # next_frontier_select_recv = [] - # for j in range(size): - # next_frontier_select_recv.append(torch.cuda.LongTensor(next_frontier_select.size())) - # dist.all_gather(next_frontier_select_recv, next_frontier_select, group) - # next_frontier_select = torch.cat(next_frontier_select_recv, dim=0) - - # batches_select_recv = [] - # for j in range(size): - # batches_select_recv.append(torch.cuda.LongTensor(batches_select.size())) - # dist.all_gather(batches_select_recv, batches_select, group) - # batches_select = torch.cat(batches_select_recv, dim=0) - print(f"rank: {rank} next_frontier_select.size: {next_frontier_select.size()}") print(f"rank: {rank} batches_select.size: {batches_select.size()}") print(f"rank: {rank} next_frontier_select: {next_frontier_select}") @@ -325,54 +313,94 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") + + bulk_row_select_rows = torch.arange(0, batches._nnz(), device=torch.device("cuda")) + bulk_row_select_cols = batches._indices()[1,:] + bulk_row_select_indices = torch.stack((bulk_row_select_rows, bulk_row_select_cols)) + bulk_row_select_values = torch.cuda.FloatTensor(batches._nnz()).fill_(1.0) + + sampled_indices, sampled_values = torch_sparse.spspmm(bulk_row_select_indices, bulk_row_select_values, + adj_matrix._indices(), adj_matrix._values(), + batch_size * mb_count, node_count_total, + node_count_total, coalesced=True) + + row_shift = torch.cuda.LongTensor(sampled_values.size(0)).fill_(0) + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), batch_size, node_count_total) + sampled_indices[1,:] += row_shift + + col_shift = torch.cuda.LongTensor(next_frontier_select.numel()).fill_(0) + bulk_col_select_rows = next_frontier_select.reshape(-1) + shift_colselect_gpu(col_shift, next_frontier_select.numel(), frontier_size + batch_size, node_count_total) + bulk_col_select_rows += col_shift + bulk_col_select_cols = torch.arange(0, next_frontier_select.numel(), device=torch.device("cuda")) + bulk_col_select_cols = torch.remainder(bulk_col_select_cols, frontier_size + batch_size) + bulk_col_select_indices = torch.stack((bulk_col_select_rows, bulk_col_select_cols)) + bulk_col_select_values = torch.cuda.FloatTensor(next_frontier_select.numel()).fill_(1.0) + + sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, + bulk_col_select_indices, bulk_col_select_values, + batch_size * mb_count, node_count_total * mb_count, + frontier_size + batch_size, coalesced=True) - torch.cuda.nvtx.range_push("nvtx-select-rowcols") for j in range(mb_count): - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-mtxs") - if i == 0: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches_select[j,:])) - else: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), \ - current_frontier[j, :])) - row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) - - col_select_mtx_indices = torch.stack((next_frontier_select[j], torch.arange(start=0, \ - end=next_frontier_select[j].size(0)).cuda())) - col_select_mtx_values = torch.cuda.FloatTensor(next_frontier_select[j].size(0)).fill_(1.0) - torch.cuda.nvtx.range_pop() - timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) - - # multiply row_select matrix with adj_matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), - row_select_mtx_values, - adj_matrix._indices(), adj_matrix._values(), - nnz, node_count, node_count, coalesced=True) - torch.cuda.nvtx.range_pop() - timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) - - # multiply adj_matrix with col_select matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, - col_select_mtx_indices.long(), col_select_mtx_values, - nnz, node_count, next_frontier_select[j].size(0), - coalesced=True) - torch.cuda.nvtx.range_pop() - timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) - # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-sample") - current_frontier[j, :] = next_frontier_select[j, :] - adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ - size=(nnz, next_frontier_select[j].size(0))) + row_select_min = j * batch_size + row_select_max = (j + 1) * batch_size + sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + adj_matrix_sample_values = sampled_values[sample_select_mask] + + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ + size=(batch_size, frontier_size + batch_size)) adj_matrices[j][i] = adj_matrix_sample - torch.cuda.nvtx.range_pop() - timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) - torch.cuda.nvtx.range_pop() + + # torch.cuda.nvtx.range_push("nvtx-select-rowcols") + # for j in range(mb_count): + # start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-select-mtxs") + # if i == 0: + # row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches_select[j,:])) + # else: + # row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), \ + # current_frontier[j, :])) + # row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) + + # col_select_mtx_indices = torch.stack((next_frontier_select[j], torch.arange(start=0, \ + # end=next_frontier_select[j].size(0)).cuda())) + # col_select_mtx_values = torch.cuda.FloatTensor(next_frontier_select[j].size(0)).fill_(1.0) + # torch.cuda.nvtx.range_pop() + # timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) + + # # multiply row_select matrix with adj_matrix + # start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") + # sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), + # row_select_mtx_values, + # adj_matrix._indices(), adj_matrix._values(), + # nnz, node_count, node_count, coalesced=True) + # torch.cuda.nvtx.range_pop() + # timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) + + # # multiply adj_matrix with col_select matrix + # start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") + # sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, + # col_select_mtx_indices.long(), col_select_mtx_values, + # nnz, node_count, next_frontier_select[j].size(0), + # coalesced=True) + # torch.cuda.nvtx.range_pop() + # timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) + # # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] + + # start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-set-sample") + # current_frontier[j, :] = next_frontier_select[j, :] + # adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ + # size=(nnz, next_frontier_select[j].size(0))) + # adj_matrices[j][i] = adj_matrix_sample + # torch.cuda.nvtx.range_pop() + # timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) + # torch.cuda.nvtx.range_pop() print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") for k, v in timing_dict.items(): diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 31e2018..3ad4d03 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -542,6 +542,53 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: CHECK_ERROR("normalize error") } +__global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, 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 = row_select_rows[i] / batch_size; + row_shift[i] += mb_row * node_count; + } +} + +void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_select_rows, + 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); + + ShiftRowSelect<<>>(row_shift.data(), + row_select_rows.data(), + nnz, + batch_size, + node_count); + 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") +} + 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"); @@ -553,4 +600,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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("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"); } From c329534e926c52d855458e30f93daf7969a44133 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 4 May 2022 18:05:04 -0700 Subject: [PATCH 054/178] Bulk row/col selection for bulkmb=8 cora --- cagnet/samplers/ladies.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index e6d6ff1..165d792 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -329,7 +329,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota sampled_indices[1,:] += row_shift col_shift = torch.cuda.LongTensor(next_frontier_select.numel()).fill_(0) - bulk_col_select_rows = next_frontier_select.reshape(-1) + bulk_col_select_rows = next_frontier_select.clone().reshape(-1) shift_colselect_gpu(col_shift, next_frontier_select.numel(), frontier_size + batch_size, node_count_total) bulk_col_select_rows += col_shift bulk_col_select_cols = torch.arange(0, next_frontier_select.numel(), device=torch.device("cuda")) @@ -348,6 +348,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ (sampled_indices[0,:] < row_select_max) adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + adj_matrix_sample_indices[0,:] -= row_select_min adj_matrix_sample_values = sampled_values[sample_select_mask] adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ From 0ef7958db95b9eeb5b276a5318fa195b7802217f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 4 May 2022 18:05:37 -0700 Subject: [PATCH 055/178] Cleanup --- cagnet/samplers/ladies.py | 48 --------------------------------------- 1 file changed, 48 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 165d792..2ac7c38 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -355,54 +355,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota size=(batch_size, frontier_size + batch_size)) adj_matrices[j][i] = adj_matrix_sample - # torch.cuda.nvtx.range_push("nvtx-select-rowcols") - # for j in range(mb_count): - # start_time(start_timer) - # torch.cuda.nvtx.range_push("nvtx-select-mtxs") - # if i == 0: - # row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches_select[j,:])) - # else: - # row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), \ - # current_frontier[j, :])) - # row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) - - # col_select_mtx_indices = torch.stack((next_frontier_select[j], torch.arange(start=0, \ - # end=next_frontier_select[j].size(0)).cuda())) - # col_select_mtx_values = torch.cuda.FloatTensor(next_frontier_select[j].size(0)).fill_(1.0) - # torch.cuda.nvtx.range_pop() - # timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) - - # # multiply row_select matrix with adj_matrix - # start_time(start_timer) - # torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") - # sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), - # row_select_mtx_values, - # adj_matrix._indices(), adj_matrix._values(), - # nnz, node_count, node_count, coalesced=True) - # torch.cuda.nvtx.range_pop() - # timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) - - # # multiply adj_matrix with col_select matrix - # start_time(start_timer) - # torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - # sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, - # col_select_mtx_indices.long(), col_select_mtx_values, - # nnz, node_count, next_frontier_select[j].size(0), - # coalesced=True) - # torch.cuda.nvtx.range_pop() - # timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) - # # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] - - # start_time(start_timer) - # torch.cuda.nvtx.range_push("nvtx-set-sample") - # current_frontier[j, :] = next_frontier_select[j, :] - # adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ - # size=(nnz, next_frontier_select[j].size(0))) - # adj_matrices[j][i] = adj_matrix_sample - # torch.cuda.nvtx.range_pop() - # timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) - # torch.cuda.nvtx.range_pop() - print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") for k, v in timing_dict.items(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") From 442f8b4462193877499aff8f0d813febf3af6802 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 4 May 2022 18:09:02 -0700 Subject: [PATCH 056/178] Cleanup --- cagnet/samplers/ladies.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 2ac7c38..f3ef9c1 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -306,13 +306,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches._values().bool()).view(mb_count, batch_size) next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) - print(f"rank: {rank} next_frontier_select.size: {next_frontier_select.size()}") - print(f"rank: {rank} batches_select.size: {batches_select.size()}") - print(f"rank: {rank} next_frontier_select: {next_frontier_select}") - print(f"rank: {rank} batches_select: {batches_select}") - torch.cuda.nvtx.range_pop() - print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") bulk_row_select_rows = torch.arange(0, batches._nnz(), device=torch.device("cuda")) bulk_row_select_cols = batches._indices()[1,:] From e245450ce8aad65a20f797dcd3c0dc2446199960 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 4 May 2022 18:21:59 -0700 Subject: [PATCH 057/178] Added timers for bulk row/col selection --- cagnet/samplers/ladies.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index f3ef9c1..38ca11a 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -291,7 +291,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) torch.cuda.nvtx.range_pop() # nvtx-sampling-iter - timing_dict["sampling_iters"].append(stop_time(sample_start_timer, sample_stop_timer)) + timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) torch.cuda.nvtx.range_pop() # nvtx-sampling @@ -305,19 +305,29 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches_select = torch.masked_select(batches._indices()[1,:], \ batches._values().bool()).view(mb_count, batch_size) next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) - torch.cuda.nvtx.range_pop() + timing_dict["construct-nextf"].append(stop_time(sample_start_timer, sample_stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-rowselect-preproc") bulk_row_select_rows = torch.arange(0, batches._nnz(), device=torch.device("cuda")) bulk_row_select_cols = batches._indices()[1,:] bulk_row_select_indices = torch.stack((bulk_row_select_rows, bulk_row_select_cols)) bulk_row_select_values = torch.cuda.FloatTensor(batches._nnz()).fill_(1.0) + torch.cuda.nvtx.range_pop() + timing_dict["rowselect-prepoc"].append(stop_time(sample_start_timer, sample_stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-rowselect-spgemm") sampled_indices, sampled_values = torch_sparse.spspmm(bulk_row_select_indices, bulk_row_select_values, adj_matrix._indices(), adj_matrix._values(), batch_size * mb_count, node_count_total, node_count_total, coalesced=True) + torch.cuda.nvtx.range_pop() + timing_dict["rowselect-spgemm"].append(stop_time(sample_start_timer, sample_stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-colselect-preproc") row_shift = torch.cuda.LongTensor(sampled_values.size(0)).fill_(0) shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), batch_size, node_count_total) sampled_indices[1,:] += row_shift @@ -330,12 +340,20 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota bulk_col_select_cols = torch.remainder(bulk_col_select_cols, frontier_size + batch_size) bulk_col_select_indices = torch.stack((bulk_col_select_rows, bulk_col_select_cols)) bulk_col_select_values = torch.cuda.FloatTensor(next_frontier_select.numel()).fill_(1.0) + torch.cuda.nvtx.range_pop() + timing_dict["colselect-prepoc"].append(stop_time(sample_start_timer, sample_stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-colselect-spgemm") sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, bulk_col_select_indices, bulk_col_select_values, batch_size * mb_count, node_count_total * mb_count, frontier_size + batch_size, coalesced=True) + torch.cuda.nvtx.range_pop() + timing_dict["colselect-spgemm"].append(stop_time(sample_start_timer, sample_stop_timer)) + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-sample-select") for j in range(mb_count): row_select_min = j * batch_size row_select_max = (j + 1) * batch_size @@ -348,6 +366,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ size=(batch_size, frontier_size + batch_size)) adj_matrices[j][i] = adj_matrix_sample + torch.cuda.nvtx.range_pop() + timing_dict["sample-select"].append(stop_time(sample_start_timer, sample_stop_timer)) print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") for k, v in timing_dict.items(): From 20655e7665db99465de840cf5b73eed4a0af0567 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 5 May 2022 12:46:03 -0700 Subject: [PATCH 058/178] Sub out spspmm with distributed spgemm --- cagnet/samplers/ladies.py | 27 ++++++++++++++++++--------- sparse-extension/sparse_coo_tensor.cu | 14 ++++++++++---- 2 files changed, 28 insertions(+), 13 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 38ca11a..534a5f5 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -319,17 +319,21 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-rowselect-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(bulk_row_select_indices, bulk_row_select_values, - adj_matrix._indices(), adj_matrix._values(), - batch_size * mb_count, node_count_total, - node_count_total, coalesced=True) + # sampled_indices, sampled_values = torch_sparse.spspmm(bulk_row_select_indices, bulk_row_select_values, + # adj_matrix._indices(), adj_matrix._values(), + # batch_size * mb_count, node_count_total, + # node_count_total, coalesced=True) + bulk_row_select = torch.sparse_coo_tensor(bulk_row_select_indices, bulk_row_select_values, + (batch_size * mb_count, node_count_total)) + sampled_indices, sampled_values = dist_spgemm1D(bulk_row_select, adj_matrix, rank, size, group) torch.cuda.nvtx.range_pop() timing_dict["rowselect-spgemm"].append(stop_time(sample_start_timer, sample_stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-colselect-preproc") row_shift = torch.cuda.LongTensor(sampled_values.size(0)).fill_(0) - shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), batch_size, node_count_total) + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, batch_size, node_count_total, mb_count_total) sampled_indices[1,:] += row_shift col_shift = torch.cuda.LongTensor(next_frontier_select.numel()).fill_(0) @@ -345,10 +349,15 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-colselect-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, - bulk_col_select_indices, bulk_col_select_values, - batch_size * mb_count, node_count_total * mb_count, - frontier_size + batch_size, coalesced=True) + # sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, + # bulk_col_select_indices, bulk_col_select_values, + # batch_size * mb_count, node_count_total * mb_count, + # frontier_size + batch_size, coalesced=True) + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + (batch_size * mb_count, node_count_total * mb_count_total)) + bulk_col_select = torch.sparse_coo_tensor(bulk_col_select_indices, bulk_col_select_values, + (node_count_total * mb_count, frontier_size + batch_size)) + sampled_indices, sampled_values = dist_spgemm1D(sample_mtx, bulk_col_select, rank, size, group) torch.cuda.nvtx.range_pop() timing_dict["colselect-spgemm"].append(stop_time(sample_start_timer, sample_stop_timer)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 3ad4d03..78abfae 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -542,18 +542,21 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: CHECK_ERROR("normalize error") } -__global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int nnz, int batch_size, int node_count) { +__global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int rank, int size, + int nnz, int batch_size, int node_count, int mb_count) { + long id = blockIdx.x * blockDim.x + threadIdx.x; long stride = blockDim.x * gridDim.x; + long proc_row_chunk = rank * ((batch_size * mb_count) / size); for (int i = id; i < nnz; i += stride) { - long mb_row = row_select_rows[i] / batch_size; + long mb_row = (row_select_rows[i] + proc_row_chunk) / batch_size; row_shift[i] += mb_row * node_count; } } void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_select_rows, - int nnz, int batch_size, int node_count) { + int nnz, int rank, int size, int batch_size, int node_count, int mb_count) { int BLOCK_SIZE = 256; @@ -562,9 +565,12 @@ void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_sele ShiftRowSelect<<>>(row_shift.data(), row_select_rows.data(), + rank, + size, nnz, batch_size, - node_count); + node_count, + mb_count); CHECK_ERROR("shift row select error") } From 4f08650ae0e50753219c9e85e4279accaedfe986 Mon Sep 17 00:00:00 2001 From: Emily Marx Date: Thu, 5 May 2022 15:00:40 -0700 Subject: [PATCH 059/178] sampled rows on diagonal; tested on cora --- cagnet/samplers/ladies.py | 111 ++++++++++++++++++++++---------------- 1 file changed, 65 insertions(+), 46 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index b5fa103..77cbdc6 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -75,7 +75,7 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = torch_sparse.spspmm(frontier_indices.long(), frontier_values, + p_num_indices, p_num_values = torch_sparse.spspmm(frontier_indices.long(), frontier_values, adj_matrix._indices().long(), adj_mat_squared, mb_count, node_count, node_count, coalesced=True) torch.cuda.nvtx.range_pop() @@ -139,9 +139,9 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - compute_darts_gpu(dartx_values, darty_values, neighbor_sizes._values(), psum_neighbor_sizes, + compute_darts_gpu(dartx_values, darty_values, neighbor_sizes._values(), psum_neighbor_sizes, max_prob, n_darts, mb_count) - + torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) @@ -213,53 +213,72 @@ def ladies_sampler(adj_matrix, batch_size, frontier_size, mb_count, n_layers, n_ print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") torch.cuda.nvtx.range_push("nvtx-select-rowcols") - for j in range(mb_count): - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-mtxs") - if i == 0: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), batches[j,:])) - else: - row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz).cuda(), \ - current_frontier[j, :])) - row_select_mtx_values = torch.cuda.FloatTensor(nnz).fill_(1.0) - - col_select_mtx_indices = torch.stack((next_frontier[j], torch.arange(start=0, \ - end=next_frontier[j].size(0)).cuda())) - col_select_mtx_values = torch.cuda.FloatTensor(next_frontier[j].size(0)).fill_(1.0) - torch.cuda.nvtx.range_pop() - timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) + # 1. Make the row/col select matrices + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-mtxs") - # multiply row_select matrix with adj_matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), - row_select_mtx_values, - adj_matrix._indices(), adj_matrix._values(), - nnz, node_count, node_count, coalesced=True) - torch.cuda.nvtx.range_pop() - timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) + if i == 0: + row_select_mtx_indices = torch.stack((torch.arange(start=0, end=nnz * mb_count).cuda(), + batches.view(-1))) + else: + # TODO i > 0 (make sure size args to spspmm are right) + print("ERROR i > 0") + + row_select_mtx_values = torch.cuda.FloatTensor(row_select_mtx_indices[0].size(0)).fill_(1.0) + + repeated_next_frontier = next_frontier.repeat_interleave(nnz, dim=0) + scale_mtx = torch.arange(start=0, end=node_count * nnz * mb_count, step=node_count).cuda() + repeated_next_frontier.add_(scale_mtx[:, None]) + col_select_mtx_rows=repeated_next_frontier.view(-1) + col_select_mtx_cols = torch.arange(next_frontier.size(1)).cuda().repeat(nnz * 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.FloatTensor(col_select_mtx_rows.size(0)).fill_(1.0) + torch.cuda.nvtx.range_pop() + timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) - # multiply adj_matrix with col_select matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, - col_select_mtx_indices.long(), col_select_mtx_values, - nnz, node_count, next_frontier[j].size(0), - coalesced=True) - torch.cuda.nvtx.range_pop() - timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) - # layer_adj_matrix = adj_matrix[current_frontier[:,0], next_frontier] + # 2. Multiply row_select matrix with adj_matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") + sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), + row_select_mtx_values, + adj_matrix._indices(), adj_matrix._values(), + nnz * mb_count, node_count, node_count, coalesced=True) + torch.cuda.nvtx.range_pop() + timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-sample") - current_frontier[j, :] = next_frontier[j, :] - adj_matrix_sample = torch.sparse_coo_tensor(indices=sampled_indices, values=sampled_values, \ - size=(nnz, next_frontier[j].size(0))) - adj_matrices[j][i] = adj_matrix_sample - torch.cuda.nvtx.range_pop() - timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) + # 3. Expand sampled rows + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-row-select-expand") + sampled_indices[1,:].add_(sampled_indices[0,:], alpha=node_count) torch.cuda.nvtx.range_pop() - + timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) + + # 4. Multiply sampled rows with col_select matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") + sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, + col_select_mtx_indices.long(), col_select_mtx_values, + nnz * mb_count, node_count * nnz * mb_count, next_frontier.size(1), + coalesced=True) + torch.cuda.nvtx.range_pop() + timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) + + # 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))) + for mb in range(mb_count): + mb_sampled_rows = torch.arange(start=batch_size * mb, end=batch_size * (mb + 1)).cuda() + adj_matrices[mb][i] = torch.index_select(adj_matrix_sample, 0, mb_sampled_rows) + + current_frontier = next_frontier + torch.cuda.nvtx.range_pop() + timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) + + torch.cuda.nvtx.range_pop() + + print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") for k, v in timing_dict.items(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") From 3934a1757a326fc4b1cef442389ecb14155f2e0d Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 5 May 2022 22:33:31 -0700 Subject: [PATCH 060/178] Moved extracting samples to main file and test file --- cagnet/samplers/ladies.py | 31 ++++++++++++++++++------------- examples/gcn_15d.py | 39 +++++++++++++++++++++++++++++++++------ tests/test_ladies.py | 5 ++++- 3 files changed, 55 insertions(+), 20 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 534a5f5..037a97f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -73,13 +73,15 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota timing_dict = defaultdict(list) - adj_matrices = [[None] * n_layers for x in range(mb_count_total)] # adj_matrices[i][j] -- mb i layer j 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) + # 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): if i == 0: @@ -363,18 +365,21 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-sample-select") - for j in range(mb_count): - row_select_min = j * batch_size - row_select_max = (j + 1) * batch_size - sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask] - - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ - size=(batch_size, frontier_size + batch_size)) - adj_matrices[j][i] = adj_matrix_sample + # for j in range(mb_count): + # row_select_min = j * batch_size + # row_select_max = (j + 1) * batch_size + # sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ + # (sampled_indices[0,:] < row_select_max) + # adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + # adj_matrix_sample_indices[0,:] -= row_select_min + # adj_matrix_sample_values = sampled_values[sample_select_mask] + + # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ + # size=(batch_size, frontier_size + batch_size)) + # adj_matrices[j][i] = adj_matrix_sample + adj_matrix_sample = torch.sparse_coo_tensor(sampled_indices, sampled_values, \ + size=(batch_size * mb_count, frontier_size + batch_size)) + adj_matrices[i] = adj_matrix_sample torch.cuda.nvtx.range_pop() timing_dict["sample-select"].append(stop_time(sample_start_timer, sample_stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 1b5dc04..bd8a6cf 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -339,7 +339,6 @@ def main(args): test_nid = data.test_mask.nonzero().squeeze() batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major - node_count = adj_matrix.size(0) torch.cuda.nvtx.range_pop() torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") @@ -374,15 +373,43 @@ def main(args): print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - rank, size, col_groups[0]) + 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, \ + rank, size, col_groups[0]) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() + # adj_matrices[i][j] -- mb i layer j + rank_n_bulkmb = int(args.n_bulkmb / size) + adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] + for i in range(args.n_layers): + for j in range(rank_n_bulkmb): + row_select_min = j * args.batch_size + row_select_max = (j + 1) * args.batch_size + + sampled_indices = adj_matrices_bulk[i]._indices() + sampled_values = adj_matrices_bulk[i]._values() + + sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask] + + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ + size=(args.batch_size, args.samp_num + args.batch_size)) + adj_matrices[j][i] = adj_matrix_sample + print(f"sample: {adj_matrices}") - return current_frontier, next_frontier, adj_matrices, g_loc # return here while testing sampling code + node_count = inputs.size(0) + 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() + adj_matrix = row_normalize(adj_matrix) + adj_matrix = torch.pow(adj_matrix, 2) + return current_frontier, next_frontier, adj_matrices, adj_matrix # return here while testing sampling code # create GCN model torch.manual_seed(0) diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 213ae0c..7ec428d 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -4,6 +4,7 @@ import numpy as np import torch +import torch.distributed as dist def main(args): """ @@ -73,7 +74,9 @@ def main(args): print("") print("Beginning LADIES sampler test") - for i in range(args.n_bulkmb): + size = dist.get_world_size() + rank_n_bulkmb = int(args.n_bulkmb / size) + 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 From 4bf5aa0c6b27ab08c736a5b2313c415473b73941 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 6 May 2022 01:01:56 -0700 Subject: [PATCH 061/178] Debugging distributed LADIES --- cagnet/samplers/ladies.py | 15 ++++------- examples/gcn_15d.py | 3 ++- sparse-extension/sparse_coo_tensor.cu | 37 +++++++++++++++------------ 3 files changed, 27 insertions(+), 28 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 037a97f..632f967 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -48,9 +48,10 @@ def dist_spgemm1D(mata, matb, rank, size, group): matb_recv_values, mata.size(0), \ chunk_col_size, matb.size(1), coalesced=True) - matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()).coalesce() matc += matc_chunk + matc = matc.coalesce() return matc._indices(), matc._values() def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ @@ -95,11 +96,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - print(f"batches: {batches}") - print(f"adj_matrices: {adj_matrix}") p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) - print(f"p_num_indices: {p_num_indices}") - print(f"p_num_values: {p_num_values}") torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -160,14 +157,15 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") - dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) + # dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) + dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") - next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).int() + 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=(mb_count, node_count_total)) @@ -297,9 +295,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() # nvtx-sampling - # TODO: Might need to downsample if > frontier_size vertices were sampled for a minibatch - print(f"next_frontier.frontier-sizes: {torch.sparse.sum(next_frontier, dim=1)}") - start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-construct-nextf") next_frontier_select = torch.masked_select(next_frontier._indices()[1,:], \ diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index bd8a6cf..172fac1 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -401,8 +401,9 @@ def main(args): size=(args.batch_size, args.samp_num + args.batch_size)) adj_matrices[j][i] = adj_matrix_sample - print(f"sample: {adj_matrices}") + # print(f"sample: {adj_matrices}") node_count = inputs.size(0) + 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)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 78abfae..15bb90a 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -31,11 +31,11 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} -__device__ long binary_searchf(float *arr, float val, long imin, long imax) { +__device__ int binary_searchf(float *arr, float val, int imin, int imax) { - long ans = -1; + int ans = 0; while (imax >= imin) { - long imid = (imin + imax) / 2; + int imid = (imin + imax) / 2; if (arr[imid] <= val) { imin = imid + 1; @@ -461,14 +461,17 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, long *h_values, int dart_count, int nnz) { +__global__ void ThrowDarts1D(float *dart_values, float *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) { - long vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz); - atomicAdd((unsigned long long *)&h_values[vtx], 1L); + int vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz); + if (vtx < 0 || vtx >= nnz) { + printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); + } + atomicAdd(&h_values[vtx], 1); } } @@ -485,29 +488,29 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, ThrowDarts1D<<>>(dart_values.data(), ps_p_values.data(), - h_values.data(), + h_values.data(), dart_count, nnz); CHECK_ERROR("dart throwing error") } -__global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, long *dart_hits_count, - long total_overflow, long nnz) { +__global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, int *dart_hits_count, + int total_overflow, int nnz) { - long id = blockIdx.x * blockDim.x + threadIdx.x; - long stride = blockDim.x * gridDim.x; + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; - for (long i = id; i < total_overflow; i += stride) { - long vtx = binary_searchf(ps_dart_hits_inv, dart_select[i], 0, nnz); - atomicAnd((unsigned long long *)&dart_hits_count[vtx], 0L); + for (int i = id; i < total_overflow; i += stride) { + int vtx = binary_searchf(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, - long total_overflow, - long nnz) { + int total_overflow, + int nnz) { int BLOCK_SIZE = 256; @@ -516,7 +519,7 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, ThrowDartsSelect<<>>(dart_select.data(), ps_dart_hits_inv.data(), - dart_hits_count.data(), + dart_hits_count.data(), total_overflow, nnz); CHECK_ERROR("selection dart throwing error") From 722b11d0ddcbf28e32b61b57f57d41cc9bdb9533 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 7 May 2022 00:25:51 -0700 Subject: [PATCH 062/178] Minor cleanup --- cagnet/samplers/ladies.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 5a13b70..0c4bbfd 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -361,9 +361,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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))) - # for mb in range(mb_count): - # mb_sampled_rows = torch.arange(start=batch_size * mb, end=batch_size * (mb + 1)).cuda() - # adj_matrices[mb][i] = torch.index_select(adj_matrix_sample, 0, mb_sampled_rows) adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier From 1f032817b876ba217ad7cc1dba1c6d4e6b584014 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 7 May 2022 01:34:03 -0700 Subject: [PATCH 063/178] Applied distributed spgemm to merged matrices code --- cagnet/samplers/ladies.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0c4bbfd..0e58ebf 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -43,6 +43,7 @@ def dist_spgemm1D(mata, matb, rank, size, group): mata_chunk_indices[1,:] -= chunk_col_start mata_chunk_values = mata._values()[chunk_col_mask] + print(f"rank: {rank} i: {i} chunk_col_size: {chunk_col_size}") matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ matb_recv_values, mata.size(0), \ @@ -319,11 +320,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota row_select_mtx_values = torch.cuda.FloatTensor(row_select_mtx_indices[0].size(0)).fill_(1.0) - repeated_next_frontier = next_frontier_select.repeat_interleave(nnz, dim=0) - scale_mtx = torch.arange(start=0, end=node_count * nnz * mb_count, step=node_count).cuda() + repeated_next_frontier = next_frontier_select.clone() + scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() repeated_next_frontier.add_(scale_mtx[:, None]) col_select_mtx_rows = repeated_next_frontier.view(-1) - col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).cuda().repeat(nnz * mb_count,1).view(-1) + col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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.FloatTensor(col_select_mtx_rows.size(0)).fill_(1.0) torch.cuda.nvtx.range_pop() @@ -332,27 +333,31 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # 2. Multiply row_select matrix with adj_matrix start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-row-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(row_select_mtx_indices.long(), - row_select_mtx_values, - adj_matrix._indices(), adj_matrix._values(), - nnz * mb_count, node_count, node_count, coalesced=True) + row_select_mtx = torch.sparse_coo_tensor(row_select_mtx_indices, row_select_mtx_values, + size=(nnz * mb_count, node_count_total)) + sampled_indices, sampled_values = dist_spgemm1D(row_select_mtx, adj_matrix, rank, size, group) + torch.cuda.nvtx.range_pop() timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) # 3. Expand sampled rows start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-row-select-expand") - sampled_indices[1,:].add_(sampled_indices[0,:], alpha=node_count) + row_shift = torch.cuda.LongTensor(sampled_values.size(0)).fill_(0) + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, batch_size, node_count_total, mb_count_total) + sampled_indices[1,:] += row_shift torch.cuda.nvtx.range_pop() timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) # 4. Multiply sampled rows with col_select matrix start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - sampled_indices, sampled_values = torch_sparse.spspmm(sampled_indices, sampled_values, - col_select_mtx_indices.long(), col_select_mtx_values, - nnz * mb_count, node_count * nnz * mb_count, next_frontier.size(1), - coalesced=True) + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + size=(nnz * mb_count, node_count_total * mb_count_total)) + 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))) + sampled_indices, sampled_values = dist_spgemm1D(sample_mtx, col_select_mtx, rank, size, group) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) From 1f0630f2958154132d43b2eccf266c50bc3de244 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sun, 8 May 2022 16:06:45 -0700 Subject: [PATCH 064/178] Computes number of darts internally --- cagnet/samplers/ladies.py | 8 +++++++- examples/gcn_15d.py | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0e58ebf..0076097 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -43,7 +43,6 @@ def dist_spgemm1D(mata, matb, rank, size, group): mata_chunk_indices[1,:] -= chunk_col_start mata_chunk_values = mata._values()[chunk_col_mask] - print(f"rank: {rank} i: {i} chunk_col_size: {chunk_col_size}") matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ matb_recv_values, mata.size(0), \ @@ -111,6 +110,13 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") + if n_darts == -1: + min_prob = torch.min(p._values()) + if min_prob.item() == 0: + print(f"probability error") + min_prob = min_prob.reciprocal() + n_darts = int(torch.sum(p._values() * min_prob).item()) + start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") next_frontier = torch.sparse_coo_tensor(indices=p._indices(), diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 172fac1..b94e8bb 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -527,7 +527,7 @@ def main(args): 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('--n-darts', default=10, type=int, + parser.add_argument('--n-darts', default=-1, type=int, help='number of darts to throw per minibatch in LADIES sampling') args = parser.parse_args() print(args) From fbe4ec9776ba9dc149f14c7888a8d85b05901414 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 10 May 2022 15:38:17 -0700 Subject: [PATCH 065/178] Added timers for spgemm and coarser n-dart computation --- cagnet/samplers/ladies.py | 44 ++++++++++++++++++++++++++++++--------- examples/gcn_15d.py | 9 +++++++- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0076097..c11d341 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -16,24 +16,40 @@ def stop_time(start_timer, stop_timer): torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) -def dist_spgemm1D(mata, matb, rank, size, group): +def stop_time_add(start_timer, stop_timer, timing_dict, range_name): + if timing_dict is not None: + timing_dict[range_name].append(stop_time(start_timer, stop_timer)) + +def dist_spgemm1D(mata, matb, rank, size, group, 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)) matc = torch.sparse_coo_tensor(size=(mata.size(0), matb.size(1))).cuda() + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-matc-inst") for i in range(size): + start_time(start_timer) matb_recv_nnz = torch.cuda.IntTensor([matb._nnz()]) dist.broadcast(matb_recv_nnz, src=i, group=group) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-bcast-nnz") + start_time(start_timer) if i == 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.FloatTensor(matb_recv_nnz.item()) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-inst-recv") + start_time(start_timer) dist.broadcast(matb_recv_indices, src=i, group=group) dist.broadcast(matb_recv_values, src=i, group=group) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-bcast-data") + start_time(start_timer) chunk_col_start = i * chunk_size chunk_col_stop = min((i + 1) * chunk_size, mata.size(1)) chunk_col_size = chunk_col_stop - chunk_col_start @@ -42,16 +58,31 @@ def dist_spgemm1D(mata, matb, rank, size, group): 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, "spgemm-preproc-local") + start_time(start_timer) 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) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-local-spgemm") + + start_time(start_timer) + matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-chunk-inst") + + start_time(start_timer) + matc_chunk = matc_chunk.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-chunk-coalesce") - matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()).coalesce() + start_time(start_timer) matc += matc_chunk + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-chunk-add") + start_time(start_timer) matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-matc-coalesce") + return matc._indices(), matc._values() def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ @@ -96,7 +127,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group) + p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group, timing_dict) torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -110,13 +141,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") - if n_darts == -1: - min_prob = torch.min(p._values()) - if min_prob.item() == 0: - print(f"probability error") - min_prob = min_prob.reciprocal() - n_darts = int(torch.sum(p._values() * min_prob).item()) - start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") next_frontier = torch.sparse_coo_tensor(indices=p._indices(), diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index b94e8bb..28605dc 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -294,6 +294,7 @@ def main(args): os.environ["MASTER_ADDR"] = args.hostname os.environ["MASTER_PORT"] = "1234" + print(f"hostname: {socket.gethostname()}", flush=True) dist.init_process_group(backend=args.dist_backend) rank = dist.get_rank() size = dist.get_world_size() @@ -362,6 +363,13 @@ def main(args): g_loc = torch.pow(g_loc, 2) torch.cuda.nvtx.range_pop() + if args.n_darts == -1: + node_count = inputs.size(0) + edge_count = adj_matrix.size(1) + avg_degree = int(edge_count / node_count) + args.n_darts = avg_degree * args.batch_size + print(f"n_darts: {args.n_darts}") + print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing @@ -402,7 +410,6 @@ def main(args): adj_matrices[j][i] = adj_matrix_sample # print(f"sample: {adj_matrices}") - node_count = inputs.size(0) adj_matrix = adj_matrix.cuda() adj_matrix = torch.sparse_coo_tensor(adj_matrix, torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0), From fd94c2bebb7e273be666a6680180d8f2ddde7572 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 10 May 2022 18:17:20 -0700 Subject: [PATCH 066/178] Moved node_count up from n_darts computation --- examples/gcn_15d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 28605dc..ae0410d 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -363,8 +363,8 @@ def main(args): g_loc = torch.pow(g_loc, 2) torch.cuda.nvtx.range_pop() + node_count = inputs.size(0) if args.n_darts == -1: - node_count = inputs.size(0) edge_count = adj_matrix.size(1) avg_degree = int(edge_count / node_count) args.n_darts = avg_degree * args.batch_size From 0fa9dc2b14ab64830e1535e274ef323d0dea61c9 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 7 Jun 2022 17:54:12 -0700 Subject: [PATCH 067/178] Preliminary 1.5D implementation with c=1 --- cagnet/samplers/ladies.py | 54 +++++++++++++++++++++++++++++++-------- examples/gcn_15d.py | 7 +++-- 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index c11d341..625f92e 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -20,23 +20,29 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if timing_dict is not None: timing_dict[range_name].append(stop_time(start_timer, stop_timer)) -def dist_spgemm1D(mata, matb, rank, size, group, timing_dict=None): +def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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)) 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, "spgemm-matc-inst") - for i in range(size): + 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=i, group=group) + dist.broadcast(matb_recv_nnz, src=q, group=col_groups[rank_col]) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-bcast-nnz") start_time(start_timer) - if i == rank: + if q == rank: matb_recv_indices = matb._indices().clone() matb_recv_values = matb._values().clone() else: @@ -45,8 +51,8 @@ def dist_spgemm1D(mata, matb, rank, size, group, timing_dict=None): stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-inst-recv") start_time(start_timer) - dist.broadcast(matb_recv_indices, src=i, group=group) - dist.broadcast(matb_recv_values, src=i, group=group) + 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, "spgemm-bcast-data") start_time(start_timer) @@ -83,10 +89,34 @@ def dist_spgemm1D(mata, matb, rank, size, group, timing_dict=None): matc = matc.coalesce() stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-matc-coalesce") + # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) + # Implement sparse allreduce with allgather + matc_nnz = torch.cuda.IntTensor([matc._nnz()]) + matc_recv_nnz = [] + for i in range(replication): + matc_recv_nnz.append(torch.cuda.IntTensor([0])) + dist.all_gather(matc_recv_nnz, matc_nnz, row_groups[rank_c]) + + matc_recv_indices = [] + matc_recv_values = [] + for i in range(replication): + matc_recv_indices.append(torch.cuda.LongTensor(2, matc_recv_nnz[i].item())) + matc_recv_values.append(torch.cuda.FloatTensor(matc_recv_nnz[i].item())) + + dist.all_gather(matc_recv_indices, matc._indices(), row_groups[rank_c]) + dist.all_gather(matc_recv_values, matc._values(), row_groups[rank_c]) + + 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 = torch.stack(matc_recv) + matc = torch.sparse.sum(matc_recv, dim=0).coalesce() + return matc._indices(), matc._values() def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ - rank, size, group): + replication, rank, size, row_groups, col_groups): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -127,7 +157,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = dist_spgemm1D(batches, adj_matrix, rank, size, group, timing_dict) + p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, row_groups, \ + col_groups, timing_dict) torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -140,6 +171,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota p_den = p_den.scatter_add_(0, p._indices()[0, :], p._values()) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") + print(f"p.nnz: {p._nnz()}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") @@ -365,7 +397,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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)) - sampled_indices, sampled_values = dist_spgemm1D(row_select_mtx, adj_matrix, rank, size, group) + sampled_indices, sampled_values = dist_spgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ + row_groups, col_groups) torch.cuda.nvtx.range_pop() timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -387,7 +420,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota size=(nnz * mb_count, node_count_total * mb_count_total)) 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))) - sampled_indices, sampled_values = dist_spgemm1D(sample_mtx, col_select_mtx, rank, size, group) + sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ + row_groups, col_groups) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index ae0410d..52fb603 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -351,6 +351,7 @@ def main(args): batches[i,:] = train_nid[idx] torch.cuda.nvtx.range_pop() + # TODO: Convert this to 1.5D partitioning batches_loc = oned_partition_mb(rank, size, batches, args.n_bulkmb) torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") @@ -376,7 +377,8 @@ def main(args): current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - rank, size, col_groups[0]) + args.replication, rank, size, \ + row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() @@ -384,7 +386,8 @@ def main(args): 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, \ - rank, size, col_groups[0]) + args.replication, rank, size, \ + row_groups, col_groups) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() From 2d0a7702eab9000a60e4451acea0383f2fd5d658 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 8 Jun 2022 17:44:58 -0700 Subject: [PATCH 068/178] Removed nnz=0 error for 1.5D spgemm --- cagnet/samplers/ladies.py | 31 +++++++++++++++------------ examples/gcn_15d.py | 11 +++++----- sparse-extension/sparse_coo_tensor.cu | 10 ++++++--- 3 files changed, 29 insertions(+), 23 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 625f92e..4615781 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -25,7 +25,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_timer = torch.cuda.Event(enable_timing=True) start_time(start_timer) - chunk_size = math.ceil(float(mata.size(1) / size)) + 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 @@ -56,8 +56,9 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-bcast-data") start_time(start_timer) - chunk_col_start = i * chunk_size - chunk_col_stop = min((i + 1) * chunk_size, mata.size(1)) + 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) @@ -90,21 +91,21 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-matc-coalesce") # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) - # Implement sparse allreduce with allgather - matc_nnz = torch.cuda.IntTensor([matc._nnz()]) - matc_recv_nnz = [] - for i in range(replication): - matc_recv_nnz.append(torch.cuda.IntTensor([0])) - dist.all_gather(matc_recv_nnz, matc_nnz, row_groups[rank_c]) + # Implement sparse allreduce w/ all_gather and padding + 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() matc_recv_indices = [] matc_recv_values = [] for i in range(replication): - matc_recv_indices.append(torch.cuda.LongTensor(2, matc_recv_nnz[i].item())) - matc_recv_values.append(torch.cuda.FloatTensor(matc_recv_nnz[i].item())) + matc_recv_indices.append(torch.cuda.LongTensor(2, matc_nnz).fill_(0)) + matc_recv_values.append(torch.cuda.FloatTensor(matc_nnz).fill_(0.0)) - dist.all_gather(matc_recv_indices, matc._indices(), row_groups[rank_c]) - dist.all_gather(matc_recv_values, matc._values(), row_groups[rank_c]) + 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.LongTensor(matc_nnz - matc._nnz()).fill_(0.0))) + 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]) matc_recv = [] for i in range(replication): @@ -399,6 +400,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota size=(nnz * mb_count, node_count_total)) sampled_indices, sampled_values = dist_spgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ row_groups, col_groups) + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + size=(nnz * mb_count, node_count_total)) torch.cuda.nvtx.range_pop() timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -408,7 +411,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-row-select-expand") row_shift = torch.cuda.LongTensor(sampled_values.size(0)).fill_(0) shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - rank, size, batch_size, node_count_total, mb_count_total) + rank, size, replication, batch_size, node_count_total, mb_count_total) sampled_indices[1,:] += row_shift torch.cuda.nvtx.range_pop() timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 52fb603..7dbf24f 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -205,10 +205,10 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep print(f"rank: {rank} inputs_loc.size: {inputs_loc.size()}", flush=True) return inputs_loc, adj_matrix_loc, am_pbyp -def oned_partition_mb(rank, size, batches, mb_count): - # TODO: Include replicaiton at some point - batch_partitions = torch.split(batches, math.ceil(float(mb_count / size)), dim=0) - return batch_partitions[rank] +def one5d_partition_mb(rank, size, batches, replication, mb_count): + rank_c = rank // replication + batch_partitions = torch.split(batches, math.ceil(float(mb_count / (size / replication))), dim=0) + return batch_partitions[rank_c] def evaluate(model, graph, features, labels, nid, nid_count, ampbyp, group): model.eval() @@ -351,8 +351,7 @@ def main(args): batches[i,:] = train_nid[idx] torch.cuda.nvtx.range_pop() - # TODO: Convert this to 1.5D partitioning - batches_loc = oned_partition_mb(rank, size, batches, args.n_bulkmb) + batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") batches_indices_rows = torch.arange(batches_loc.size(0)).cuda() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 15bb90a..75c999e 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -546,12 +546,14 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: } __global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int rank, int size, - int nnz, int batch_size, int node_count, int mb_count) { + int replication, int nnz, int batch_size, int node_count, int mb_count) { long id = blockIdx.x * blockDim.x + threadIdx.x; long stride = blockDim.x * gridDim.x; - long proc_row_chunk = rank * ((batch_size * mb_count) / size); + 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) / batch_size; row_shift[i] += mb_row * node_count; @@ -559,7 +561,8 @@ __global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int rank, } void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_select_rows, - int nnz, int rank, int size, int batch_size, int node_count, int mb_count) { + int nnz, int rank, int size, int replication, int batch_size, int node_count, + int mb_count) { int BLOCK_SIZE = 256; @@ -570,6 +573,7 @@ void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_sele row_select_rows.data(), rank, size, + replication, nnz, batch_size, node_count, From 614e58d6643e52f4a1e6b7d6d182b7dbffd94eee Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 9 Jun 2022 18:07:26 -0700 Subject: [PATCH 069/178] Fixed stray zero issue --- cagnet/samplers/ladies.py | 2 ++ examples/gcn_15d.py | 6 +++++- tests/test_ladies.py | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 4615781..01ce13b 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -59,6 +59,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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)) + print(f"chunk_col_start: {chunk_col_start} chunk_col_stop: {chunk_col_stop}") chunk_col_size = chunk_col_stop - chunk_col_start chunk_col_mask = (mata._indices()[1, :] >= chunk_col_start) & (mata._indices()[1, :] < chunk_col_stop) @@ -433,6 +434,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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))) + print(f"adj_matrix_sample: {adj_matrix_sample}") adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 7dbf24f..60a2e3b 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -391,7 +391,7 @@ def main(args): torch.cuda.profiler.cudart().cudaProfilerStop() # adj_matrices[i][j] -- mb i layer j - rank_n_bulkmb = int(args.n_bulkmb / size) + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] for i in range(args.n_layers): for j in range(rank_n_bulkmb): @@ -401,6 +401,10 @@ def main(args): sampled_indices = adj_matrices_bulk[i]._indices() sampled_values = adj_matrices_bulk[i]._values() + nnz_mask = sampled_values != 0 + sampled_indices = sampled_indices[:, nnz_mask] + sampled_values = sampled_values[nnz_mask] + sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ (sampled_indices[0,:] < row_select_max) adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 7ec428d..89a3f43 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -75,7 +75,7 @@ def main(args): print("") print("Beginning LADIES sampler test") size = dist.get_world_size() - rank_n_bulkmb = int(args.n_bulkmb / 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 From fe7fa7cf7668ddf88ec2671ebedc9a01044b5b1f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 10 Jun 2022 12:52:03 -0700 Subject: [PATCH 070/178] Moved unmasking from test script to spgemm func --- cagnet/samplers/ladies.py | 9 +++++++-- examples/gcn_15d.py | 4 ---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 01ce13b..08b87a0 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -59,7 +59,6 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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)) - print(f"chunk_col_start: {chunk_col_start} chunk_col_stop: {chunk_col_stop}") chunk_col_size = chunk_col_stop - chunk_col_start chunk_col_mask = (mata._indices()[1, :] >= chunk_col_start) & (mata._indices()[1, :] < chunk_col_stop) @@ -115,7 +114,12 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, matc_recv = torch.stack(matc_recv) matc = torch.sparse.sum(matc_recv, dim=0).coalesce() - return matc._indices(), matc._values() + nnz_mask = matc._values() != 0 + + matc_nnz_indices = matc._indices()[:, nnz_mask] + matc_nnz_values = matc._values()[nnz_mask] + + return matc_nnz_indices, matc_nnz_values def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ replication, rank, size, row_groups, col_groups): @@ -367,6 +371,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches_select = torch.masked_select(batches._indices()[1,:], \ batches._values().bool()).view(mb_count, batch_size) next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) + print(f"next_frontier_select: {next_frontier_select}") torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 60a2e3b..65a1b48 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -401,10 +401,6 @@ def main(args): sampled_indices = adj_matrices_bulk[i]._indices() sampled_values = adj_matrices_bulk[i]._values() - nnz_mask = sampled_values != 0 - sampled_indices = sampled_indices[:, nnz_mask] - sampled_values = sampled_values[nnz_mask] - sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ (sampled_indices[0,:] < row_select_max) adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] From 19f92b0aa758317b358b0f4a97073b9ad3c7c088 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 14 Jun 2022 00:16:19 -0700 Subject: [PATCH 071/178] Precision necessary -- convert from float to double --- cagnet/samplers/ladies.py | 75 ++-- examples/gcn_15d.py | 41 ++- examples/tmp.out | 263 ++++++++++++++ examples/tmp2.out | 312 +++++++++++++++++ examples/tmp3.out | 473 ++++++++++++++++++++++++++ examples/tmp_bsearch_fpoint_error.out | 259 ++++++++++++++ sparse-extension/sparse_coo_tensor.cu | 75 ++-- 7 files changed, 1442 insertions(+), 56 deletions(-) create mode 100644 examples/tmp.out create mode 100644 examples/tmp2.out create mode 100644 examples/tmp3.out create mode 100644 examples/tmp_bsearch_fpoint_error.out diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 08b87a0..d4fdc2c 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -47,7 +47,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, matb_recv_values = matb._values().clone() else: matb_recv_indices = torch.cuda.LongTensor(2, matb_recv_nnz.item()) - matb_recv_values = torch.cuda.FloatTensor(matb_recv_nnz.item()) + matb_recv_values = torch.cuda.DoubleTensor(matb_recv_nnz.item()) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-inst-recv") start_time(start_timer) @@ -92,32 +92,44 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, # 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, "spgemm-reduce-nnz") + 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.FloatTensor(matc_nnz).fill_(0.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.LongTensor(matc_nnz - matc._nnz()).fill_(0.0))) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-padding") + + 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, "spgemm-allgather") + 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())) + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-preproc-reduce") + start_time(start_timer) matc_recv = torch.stack(matc_recv) matc = torch.sparse.sum(matc_recv, dim=0).coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-reduce") + 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, "spgemm-unpad") return matc_nnz_indices, matc_nnz_values @@ -170,11 +182,13 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") + p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) + p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) - p_den = torch.cuda.FloatTensor(mb_count).fill_(0) - p_den = p_den.scatter_add_(0, p._indices()[0, :], p._values()) + print(f"before p._values(): {p._values()}") + # print(f"before p._values()[56912]: {p._values()[56912]} p_den[row[56912]]: {p_den[p._indices()[0, 56912]]}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") print(f"p.nnz: {p._nnz()}") @@ -192,10 +206,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota iter_count = 0 selection_iter_count = 0 torch.cuda.nvtx.range_push("nvtx-sampling") - p_rowsum = torch.cuda.FloatTensor(mb_count).fill_(0) # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True + return_dart_counts = None + return_dart_map = None + return_dart_val = None + return_ps_pvals = None + while underfull_minibatches: iter_count += 1 start_time(sample_start_timer) @@ -203,24 +221,23 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-prob-rowsum") - 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 - ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) + ps_p_values = torch.cumsum(p._values().double(), dim=0).roll(1) ps_p_values[0] = 0 + # print(f"p._values()[56912]: {p._values()[56912]} ps_p_values[56912]: {ps_p_values[56912]}") torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") - dart_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() + dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - compute_darts1d_gpu(dart_values, p_rowsum, ps_p_rowsum, n_darts, mb_count) + print(f"iter_count: {iter_count} dart_values: {dart_values}") + compute_darts1d_gpu(dart_values, n_darts, mb_count) + print(f"iter_count: {iter_count} dart_values_scaled: {dart_values}") torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) @@ -228,7 +245,21 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-filter-darts") # dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) - throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, n_darts * mb_count, p._nnz()) + dart_hits_map = torch.cuda.IntTensor(n_darts * mb_count) + throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, dart_hits_map, \ + n_darts * mb_count, p._nnz()) + if iter_count == 1: + return_dart_counts = dart_hits_count.clone() + return_dart_map = dart_hits_map.clone() + return_dart_val = dart_values.clone() + return_ps_pvals = ps_p_values.clone() + + print(f"iter_count: {iter_count} dart_hits_count.nnzidxs: {dart_hits_count.nonzero()}") + print(f"iter_count: {iter_count} dart_hits_count.nnzidxs.size: {dart_hits_count.nonzero().size()}") + print(f"iter_count: {iter_count} dart_hits_count.sum: {torch.sum(dart_hits_count)}") + print(f"iter_count: {iter_count} dart_hits_map: {dart_hits_map}") + print(f"iter_count: {iter_count} dart_hits_count: {dart_hits_count[dart_hits_count.nonzero()]}") + print(f"iter_count: {iter_count} ps_p_values: {ps_p_values}", flush=True) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) @@ -271,7 +302,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-reciprocal") - dart_hits_inv = dart_hits_count.reciprocal() + 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)) @@ -285,7 +316,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-invsum") # dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() - dart_hits_inv_sum = torch.cuda.FloatTensor(mb_count).fill_(0) + dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) @@ -297,7 +328,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-computedarts") - dart_select = torch.cuda.FloatTensor(total_overflow).uniform_() + dart_select = torch.cuda.DoubleTensor(total_overflow).uniform_() # Compute darts for selection compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) @@ -371,7 +402,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches_select = torch.masked_select(batches._indices()[1,:], \ batches._values().bool()).view(mb_count, batch_size) next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) - print(f"next_frontier_select: {next_frontier_select}") torch.cuda.nvtx.range_pop() print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") @@ -387,7 +417,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO i > 0 (make sure size args to spspmm are right) print("ERROR i > 0") - row_select_mtx_values = torch.cuda.FloatTensor(row_select_mtx_indices[0].size(0)).fill_(1.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() scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() @@ -395,7 +425,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota col_select_mtx_rows = repeated_next_frontier.view(-1) col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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.FloatTensor(col_select_mtx_rows.size(0)).fill_(1.0) + col_select_mtx_values = torch.cuda.DoubleTensor(col_select_mtx_rows.size(0)).fill_(1.0) torch.cuda.nvtx.range_pop() timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) @@ -439,7 +469,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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))) - print(f"adj_matrix_sample: {adj_matrix_sample}") adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier @@ -448,9 +477,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() - print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}") + print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) for k, v in timing_dict.items(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") - return batches_select, next_frontier_select, adj_matrices + return batches_select, next_frontier_select, adj_matrices, return_dart_counts, return_dart_map, return_dart_val, return_ps_pvals diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 65a1b48..c2cf046 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -190,7 +190,7 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep 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)), + 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], \ @@ -325,7 +325,7 @@ def main(args): print("coalescing", flush=True) g_loc = g_loc.coalesce().t() print("normalizing", flush=True) - g_loc = row_normalize(g_loc) + # g_loc = row_normalize(g_loc) print("done normalizing", flush=True) n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) @@ -358,7 +358,7 @@ def main(args): 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_values = torch.cuda.DoubleTensor(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) torch.cuda.nvtx.range_pop() @@ -373,7 +373,9 @@ def main(args): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - current_frontier, next_frontier, adj_matrices = ladies_sampler(g_loc, batches_loc, args.batch_size, \ + torch.manual_seed(0) + current_frontier, next_frontier, adj_matrices, dart_count1, dart_map1, dart_vals1, ps_pvals1 = \ + ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ args.replication, rank, size, \ @@ -382,11 +384,40 @@ def main(args): print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - current_frontier, next_frontier, adj_matrices_bulk = ladies_sampler(g_loc, batches_loc, args.batch_size, \ + torch.manual_seed(0) + current_frontier, next_frontier, adj_matrices_bulk, dart_count2, dart_map2, dart_vals2, ps_pvals2 = \ + ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ args.replication, rank, size, \ row_groups, col_groups) + print() + print(f"return_dart_count1.nnzidxs: {dart_count1.nonzero().squeeze()}") + print(f"return_dart_count1: {dart_count1[dart_count1.nonzero().squeeze()]}") + print(f"return_dart_map1: {dart_map1}") + print(f"return_dart_val1: {dart_vals1}") + print(f"return_dart_count2.nnzidxs: {dart_count2.nonzero().squeeze()}") + print(f"return_dart_count2: {dart_count2[dart_count2.nonzero().squeeze()]}") + print(f"return_dart_map2: {dart_map2}") + print(f"return_dart_val2: {dart_vals2}") + print() + unmatched_darts = ((dart_map1 == dart_map2) == 0).nonzero().squeeze() + print(f"unmatched_darts: {unmatched_darts}") + # print(f"unmatched_darts.size: {unmatched_darts.size()}") + # print(f"dart_map1[unmatched]: {dart_map1[unmatched_darts]}") + # print(f"dart_map2[unmatched]: {dart_map2[unmatched_darts]}") + # print() + # print(f"dart_vals1[unmatched]: {dart_vals1[unmatched_darts]}") + # print(f"dart_vals2[unmatched]: {dart_vals2[unmatched_darts]}") + # print() + # print(f"ps_pvasl1.size: {ps_pvals1.size()}") + # print(f"ps_pvals1: {ps_pvals1}") + # print(f"ps_pvals2: {ps_pvals2}") + # print(f"ps_pvals1[map1unmatched]: {ps_pvals1[dart_map1[unmatched_darts].long()]}") + # print(f"ps_pvals2[map1unmatched]: {ps_pvals2[dart_map1[unmatched_darts].long()]}") + # print() + # print(f"ps_pvals1[map2unmatched]: {ps_pvals1[dart_map2[unmatched_darts].long()]}") + # print(f"ps_pvals2[map2unmatched]: {ps_pvals2[dart_map2[unmatched_darts].long()]}") torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() diff --git a/examples/tmp.out b/examples/tmp.out new file mode 100644 index 0000000..9698812 --- /dev/null +++ b/examples/tmp.out @@ -0,0 +1,263 @@ +0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=512, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) +0: hostname: nid002964 +0: hostname: nid002964 rank: 0 size: 1 +0: start partitioning +0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) +0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) +0: end partitioning +0: coalescing +0: normalizing +0: done normalizing +0: imid: 18166615 val: 0.817233 arr[imid]: 256.417358 +0: imid: 9083307 val: 0.817233 arr[imid]: 128.898560 +0: imid: 4541653 val: 0.817233 arr[imid]: 64.081360 +0: imid: 2270826 val: 0.817233 arr[imid]: 31.746990 +0: imid: 1135412 val: 0.817233 arr[imid]: 15.978256 +0: imid: 567705 val: 0.817233 arr[imid]: 8.033811 +0: imid: 283852 val: 0.817233 arr[imid]: 4.154286 +0: imid: 141925 val: 0.817233 arr[imid]: 2.080449 +0: imid: 70962 val: 0.817233 arr[imid]: 0.972775 +0: imid: 35480 val: 0.817233 arr[imid]: 0.488077 +0: imid: 53221 val: 0.817233 arr[imid]: 0.730208 +0: imid: 62091 val: 0.817233 arr[imid]: 0.852156 +0: imid: 57656 val: 0.817233 arr[imid]: 0.791403 +0: imid: 59873 val: 0.817233 arr[imid]: 0.820989 +0: imid: 58764 val: 0.817233 arr[imid]: 0.806108 +0: imid: 59318 val: 0.817233 arr[imid]: 0.813626 +0: imid: 59595 val: 0.817233 arr[imid]: 0.817638 +0: imid: 59456 val: 0.817233 arr[imid]: 0.815839 +0: imid: 59525 val: 0.817233 arr[imid]: 0.816724 +0: imid: 59560 val: 0.817233 arr[imid]: 0.817125 +0: imid: 59577 val: 0.817233 arr[imid]: 0.817343 +0: imid: 59568 val: 0.817233 arr[imid]: 0.817230 +0: imid: 59572 val: 0.817233 arr[imid]: 0.817272 +0: imid: 59570 val: 0.817233 arr[imid]: 0.817251 +0: imid: 59569 val: 0.817233 arr[imid]: 0.817237 +0: imid: 18166615 val: 0.817233 arr[imid]: 256.417145 +0: imid: 9083307 val: 0.817233 arr[imid]: 128.898499 +0: imid: 4541653 val: 0.817233 arr[imid]: 64.081367 +0: imid: 2270826 val: 0.817233 arr[imid]: 31.746983 +0: imid: 1135412 val: 0.817233 arr[imid]: 15.978255 +0: imid: 567705 val: 0.817233 arr[imid]: 8.033811 +0: imid: 283852 val: 0.817233 arr[imid]: 4.154286 +0: imid: 141925 val: 0.817233 arr[imid]: 2.080449 +0: imid: 70962 val: 0.817233 arr[imid]: 0.972775 +0: imid: 35480 val: 0.817233 arr[imid]: 0.488077 +0: imid: 53221 val: 0.817233 arr[imid]: 0.730208 +0: imid: 62091 val: 0.817233 arr[imid]: 0.852156 +0: imid: 57656 val: 0.817233 arr[imid]: 0.791403 +0: imid: 59873 val: 0.817233 arr[imid]: 0.820989 +0: imid: 58764 val: 0.817233 arr[imid]: 0.806108 +0: imid: 59318 val: 0.817233 arr[imid]: 0.813626 +0: imid: 59595 val: 0.817233 arr[imid]: 0.817638 +0: imid: 59456 val: 0.817233 arr[imid]: 0.815839 +0: imid: 59525 val: 0.817233 arr[imid]: 0.816724 +0: imid: 59560 val: 0.817233 arr[imid]: 0.817125 +0: imid: 59577 val: 0.817233 arr[imid]: 0.817343 +0: imid: 59568 val: 0.817233 arr[imid]: 0.817230 +0: imid: 59572 val: 0.817233 arr[imid]: 0.817272 +0: imid: 59570 val: 0.817233 arr[imid]: 0.817251 +0: imid: 59569 val: 0.817233 arr[imid]: 0.817237 +0: n_darts: 125952 +0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], +0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), +0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), +0: device='cuda:0', size=(232965, 232965), nnz=114848857, +0: layout=torch.sparse_coo) +0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 511, 511, 511], +0: [195991, 117560, 100852, ..., 210813, 53334, 131122]]), +0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), +0: device='cuda:0', size=(512, 232965), nnz=131072, layout=torch.sparse_coo) +0: probability-spgemm: 473.3919677734375 +0: before p._values(): tensor([3., 1., 3., ..., 1., 1., 1.], device='cuda:0') +0: before p._values()[56912]: 2.0 p_den[row[56912]]: 142332.0 +0: compute-p: 61.10249710083008 +0: p.nnz: 36333230 +0: pre-loop: 1.451200008392334 +0: p._values()[56912]: 1.4051654034119565e-05 ps_p_values[56912]: 0.7808363437652588 +0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') +0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 1], +0: [ 3], +0: [ 4], +0: ..., +0: [36333219], +0: [36333221], +0: [36333226]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([18000704, 1]) +0: iter_count: 1 dart_hits_count.sum: 64487424 +0: iter_count: 1 dart_hits_map: tensor([ 28897, 37545, 1807, ..., 36289331, 36332750, 36304529], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[3], +0: [1], +0: [1], +0: ..., +0: [2], +0: [3], +0: [2]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 2.1077e-05, 2.8103e-05, ..., 5.1200e+02, 5.1200e+02, +0: 5.1200e+02], device='cuda:0') +0: construct-nextf: 0.834559977054596 +0: total_time: 1820.016845703125 +0: spgemm-matc-inst total_time: 0.060736000537872314 avg_time 0.060736000537872314 +0: spgemm-bcast-nnz total_time: 123.55779266357422 avg_time 123.55779266357422 +0: spgemm-inst-recv total_time: 3.5558080673217773 avg_time 3.5558080673217773 +0: spgemm-bcast-data total_time: 0.09337600320577621 avg_time 0.09337600320577621 +0: spgemm-preproc-local total_time: 0.5993279814720154 avg_time 0.5993279814720154 +0: spgemm-local-spgemm total_time: 223.27728271484375 avg_time 223.27728271484375 +0: spgemm-chunk-inst total_time: 1.167680025100708 avg_time 1.167680025100708 +0: spgemm-chunk-coalesce total_time: 33.264705657958984 avg_time 33.264705657958984 +0: spgemm-chunk-add total_time: 1.1167999505996704 avg_time 1.1167999505996704 +0: spgemm-matc-coalesce total_time: 0.01235199999064207 avg_time 0.01235199999064207 +0: spgemm-reduce-nnz total_time: 37.626495361328125 avg_time 37.626495361328125 +0: spgemm-padding total_time: 1.8718719482421875 avg_time 1.8718719482421875 +0: spgemm-allgather total_time: 2.272831916809082 avg_time 2.272831916809082 +0: spgemm-preproc-reduce total_time: 1.1596800088882446 avg_time 1.1596800088882446 +0: spgemm-reduce total_time: 40.225215911865234 avg_time 40.225215911865234 +0: spgemm-unpad total_time: 3.052704095840454 avg_time 3.052704095840454 +0: prob-rowsum total_time: 0.612608015537262 avg_time 0.612608015537262 +0: gen-darts total_time: 0.23974399268627167 avg_time 0.23974399268627167 +0: dart-throw total_time: 1.302880048751831 avg_time 1.302880048751831 +0: filter_darts total_time: 9.320735931396484 avg_time 9.320735931396484 +0: add-to-frontier total_time: 1.856287956237793 avg_time 1.856287956237793 +0: count_samples total_time: 30.434911727905273 avg_time 30.434911727905273 +0: select-preproc total_time: 0.15148800611495972 avg_time 0.15148800611495972 +0: select-psoverflow total_time: 0.6426879912614822 avg_time 0.06426879912614822 +0: select-reciprocal total_time: 13.074944019317627 avg_time 1.3074944019317627 +0: select-instmtx total_time: 11.400192022323608 avg_time 1.1400192022323608 +0: select-invsum total_time: 589.2849617004395 avg_time 58.92849617004394 +0: select-psinv total_time: 0.7217599973082542 avg_time 0.07217599973082542 +0: select-computedarts total_time: 1.292479995638132 avg_time 0.12924799956381322 +0: select-throwdarts total_time: 5.817088067531586 avg_time 0.5817088067531586 +0: select-add-to-frontier total_time: 16.777632117271423 avg_time 1.6777632117271424 +0: select-samplecount total_time: 307.49686431884766 avg_time 30.749686431884765 +0: select-overflow total_time: 0.7064639925956726 avg_time 0.07064639925956726 +0: select-iter total_time: 949.5877990722656 avg_time 94.95877990722656 +0: dart-selection total_time: 949.9708251953125 avg_time 949.9708251953125 +0: set-probs total_time: 1.7460800409317017 avg_time 1.7460800409317017 +0: compute-bool total_time: 0.056543998420238495 avg_time 0.056543998420238495 +0: sampling-iters total_time: 997.7376098632812 avg_time 997.7376098632812 +0: select-mtxs total_time: 1.0613759756088257 avg_time 1.0613759756088257 +0: row-select-spgemm total_time: 201.4571533203125 avg_time 201.4571533203125 +0: row-select-expand total_time: 2.8359999656677246 avg_time 2.8359999656677246 +0: col-select-spgemm total_time: 79.69513702392578 avg_time 79.69513702392578 +0: set-sample total_time: 0.1319040060043335 avg_time 0.1319040060043335 +0: iter_count: 1 +0: selection_iter_count: 10 +0: +0: probability-spgemm: 184.28134155273438 +0: before p._values(): tensor([3., 1., 3., ..., 1., 1., 1.], device='cuda:0') +0: before p._values()[56912]: 2.0 p_den[row[56912]]: 142332.0 +0: compute-p: 60.85647964477539 +0: p.nnz: 36333230 +0: pre-loop: 1.4413119554519653 +0: p._values()[56912]: 1.4051654034119565e-05 ps_p_values[56912]: 0.7808363437652588 +0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') +0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 0], +0: [ 1], +0: [ 3], +0: ..., +0: [36333219], +0: [36333221], +0: [36333224]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([17999426, 1]) +0: iter_count: 1 dart_hits_count.sum: 64487424 +0: iter_count: 1 dart_hits_map: tensor([ 28897, 37545, 1807, ..., 36289345, 36332774, 36304546], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[36], +0: [ 3], +0: [ 1], +0: ..., +0: [ 4], +0: [ 3], +0: [ 2]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 2.1077e-05, 2.8103e-05, ..., 5.1200e+02, 5.1200e+02, +0: 5.1200e+02], device='cuda:0') +0: construct-nextf: 0.7904639840126038 +0: total_time: 1524.7774658203125 +0: spgemm-matc-inst total_time: 0.04438399896025658 avg_time 0.04438399896025658 +0: spgemm-bcast-nnz total_time: 0.08675199747085571 avg_time 0.08675199747085571 +0: spgemm-inst-recv total_time: 3.5789759159088135 avg_time 3.5789759159088135 +0: spgemm-bcast-data total_time: 0.06678400188684464 avg_time 0.06678400188684464 +0: spgemm-preproc-local total_time: 0.22096000611782074 avg_time 0.22096000611782074 +0: spgemm-local-spgemm total_time: 96.86115264892578 avg_time 96.86115264892578 +0: spgemm-chunk-inst total_time: 1.154528021812439 avg_time 1.154528021812439 +0: spgemm-chunk-coalesce total_time: 33.27449417114258 avg_time 33.27449417114258 +0: spgemm-chunk-add total_time: 1.1093759536743164 avg_time 1.1093759536743164 +0: spgemm-matc-coalesce total_time: 0.012256000190973282 avg_time 0.012256000190973282 +0: spgemm-reduce-nnz total_time: 0.09273599833250046 avg_time 0.09273599833250046 +0: spgemm-padding total_time: 1.8496320247650146 avg_time 1.8496320247650146 +0: spgemm-allgather total_time: 2.2429440021514893 avg_time 2.2429440021514893 +0: spgemm-preproc-reduce total_time: 1.153056025505066 avg_time 1.153056025505066 +0: spgemm-reduce total_time: 39.10028839111328 avg_time 39.10028839111328 +0: spgemm-unpad total_time: 3.0185279846191406 avg_time 3.0185279846191406 +0: prob-rowsum total_time: 0.5772160291671753 avg_time 0.5772160291671753 +0: gen-darts total_time: 0.21264000236988068 avg_time 0.21264000236988068 +0: dart-throw total_time: 1.2690880298614502 avg_time 1.2690880298614502 +0: filter_darts total_time: 9.313535690307617 avg_time 9.313535690307617 +0: add-to-frontier total_time: 1.840448021888733 avg_time 1.840448021888733 +0: count_samples total_time: 30.42483139038086 avg_time 30.42483139038086 +0: select-preproc total_time: 0.07369600236415863 avg_time 0.07369600236415863 +0: select-psoverflow total_time: 0.6427519991993904 avg_time 0.06427519991993905 +0: select-reciprocal total_time: 13.018144130706787 avg_time 1.3018144130706788 +0: select-instmtx total_time: 11.403391718864441 avg_time 1.1403391718864442 +0: select-invsum total_time: 589.0520629882812 avg_time 58.905206298828126 +0: select-psinv total_time: 0.7290239855647087 avg_time 0.07290239855647088 +0: select-computedarts total_time: 1.2980800233781338 avg_time 0.12980800233781337 +0: select-throwdarts total_time: 5.81113600730896 avg_time 0.581113600730896 +0: select-add-to-frontier total_time: 16.790783882141113 avg_time 1.6790783882141114 +0: select-samplecount total_time: 307.7421760559082 avg_time 30.77421760559082 +0: select-overflow total_time: 0.7125119939446449 avg_time 0.07125119939446449 +0: select-iter total_time: 949.5481796264648 avg_time 94.95481796264649 +0: dart-selection total_time: 949.9288330078125 avg_time 949.9288330078125 +0: set-probs total_time: 1.6976959705352783 avg_time 1.6976959705352783 +0: compute-bool total_time: 0.045791998505592346 avg_time 0.045791998505592346 +0: sampling-iters total_time: 997.4190063476562 avg_time 997.4190063476562 +0: select-mtxs total_time: 0.9089279770851135 avg_time 0.9089279770851135 +0: row-select-spgemm total_time: 196.4358367919922 avg_time 196.4358367919922 +0: row-select-expand total_time: 2.838144063949585 avg_time 2.838144063949585 +0: col-select-spgemm total_time: 79.38406372070312 avg_time 79.38406372070312 +0: set-sample total_time: 0.12310399860143661 avg_time 0.12310399860143661 +0: iter_count: 1 +0: selection_iter_count: 10 +0: +0: return_dart_count1.nnzidxs: tensor([ 1, 3, 4, ..., 36333219, 36333221, 36333226], +0: device='cuda:0') +0: return_dart_count1: tensor([3, 1, 1, ..., 2, 3, 2], device='cuda:0', dtype=torch.int32) +0: return_dart_map1: tensor([ 28897, 37545, 1807, ..., 36289331, 36332750, 36304529], +0: device='cuda:0', dtype=torch.int32) +0: return_dart_val1: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: return_dart_count2.nnzidxs: tensor([ 0, 1, 3, ..., 36333219, 36333221, 36333224], +0: device='cuda:0') +0: return_dart_count2: tensor([36, 3, 1, ..., 4, 3, 2], device='cuda:0', dtype=torch.int32) +0: return_dart_map2: tensor([ 28897, 37545, 1807, ..., 36289345, 36332774, 36304546], +0: device='cuda:0', dtype=torch.int32) +0: return_dart_val2: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: +0: unmatched_darts: tensor([ 6396, 6582, 7501, ..., 64487421, 64487422, 64487423], +0: device='cuda:0') +0: unmatched_darts.size: torch.Size([53266787]) +0: dart_map1[unmatched]: tensor([ 68599, 58618, 35396, ..., 36289331, 36332750, 36304529], +0: device='cuda:0', dtype=torch.int32) +0: dart_map2[unmatched]: tensor([ 68598, 58619, 35397, ..., 36289345, 36332774, 36304546], +0: device='cuda:0', dtype=torch.int32) +0: +0: dart_vals1[unmatched]: tensor([9.4084e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: dart_vals2[unmatched]: tensor([9.4084e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: +0: ps_pvals1[map1unmatched]: tensor([9.4086e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: ps_pvals2[map1unmatched]: tensor([9.4086e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: +0: ps_pvals1[map2unmatched]: tensor([9.4084e-01, 8.0436e-01, 4.8711e-01, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: ps_pvals2[map2unmatched]: tensor([9.4084e-01, 8.0436e-01, 4.8711e-01, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') diff --git a/examples/tmp2.out b/examples/tmp2.out new file mode 100644 index 0000000..aea3f83 --- /dev/null +++ b/examples/tmp2.out @@ -0,0 +1,312 @@ +0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=1, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) +0: hostname: nid002964 +0: hostname: nid002964 rank: 0 size: 1 +0: start partitioning +0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) +0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) +0: end partitioning +0: coalescing +0: normalizing +0: done normalizing +0: n_darts: 125952 +0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], +0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), +0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), +0: device='cuda:0', size=(232965, 232965), nnz=114848857, +0: dtype=torch.float64, layout=torch.sparse_coo) +0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0, 0, 0, 0, +0: 0, 0, 0, 0], +0: [195991, 117560, 100852, 197262, 114239, 24469, 97172, +0: 12309, 212877, 77891, 135178, 163418, 60833, 125788, +0: 122201, 165750, 106560, 62024, 203033, 198683, 118941, +0: 2313, 90070, 217127, 153723, 220936, 106445, 76218, +0: 201129, 106587, 110028, 193334, 20554, 79072, 127722, +0: 137651, 214143, 177731, 30765, 91760, 101839, 39926, +0: 194008, 9498, 190477, 212081, 179417, 135134, 203107, +0: 230411, 8706, 185588, 52276, 184932, 61362, 95713, +0: 80064, 175943, 221178, 119857, 129676, 10259, 174825, +0: 208831, 61946, 99332, 214464, 184739, 49674, 164221, +0: 59486, 177651, 77796, 85412, 13951, 48978, 62602, +0: 225615, 182931, 200535, 207967, 103437, 70222, 223332, +0: 85366, 157946, 192683, 23053, 92330, 197448, 15910, +0: 19101, 171437, 179931, 52280, 197172, 60523, 154098, +0: 186386, 201681, 193268, 80456, 201243, 11949, 33508, +0: 19455, 115980, 46319, 39589, 170252, 174478, 97713, +0: 163749, 148905, 144319, 11441, 54393, 61999, 100190, +0: 83454, 114559, 59890, 39951, 88308, 88733, 142830, +0: 72414, 32648, 114384, 126907, 96293, 197589, 165314, +0: 55684, 120097, 134402, 59885, 114336, 113777, 36989, +0: 192085, 169890, 5813, 31575, 230457, 20423, 107725, +0: 77745, 13674, 13136, 12013, 213858, 97307, 180094, +0: 5329, 9922, 55727, 83586, 221433, 133085, 125773, +0: 209438, 151068, 79892, 137760, 232603, 57682, 110255, +0: 37831, 86141, 56670, 87840, 146715, 147534, 150393, +0: 58304, 206360, 56253, 196089, 48612, 28102, 53804, +0: 207564, 161368, 192228, 24247, 84217, 120440, 54570, +0: 159631, 93878, 1928, 169073, 122680, 135211, 6105, +0: 12968, 66029, 142093, 149634, 178465, 34053, 58560, +0: 161657, 130137, 17574, 137016, 80219, 176259, 101294, +0: 88046, 42065, 199930, 120417, 75237, 71001, 114135, +0: 110924, 20081, 207091, 64444, 105467, 163630, 195228, +0: 7255, 139108, 89413, 227103, 94165, 56486, 228263, +0: 223212, 214386, 32605, 75837, 73859, 67370, 125235, +0: 8321, 116827, 168590, 8279, 152684, 145955, 26546, +0: 137981, 13286, 77317, 25850, 83672, 84947, 155106, +0: 36568, 15526, 216693, 74749]]), +0: values=tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., +0: 1., 1., 1., 1.]), +0: device='cuda:0', size=(1, 232965), nnz=256, dtype=torch.float64, +0: layout=torch.sparse_coo) +0: probability-spgemm: 339.2951354980469 +0: before p._values(): tensor([ 7.8125e-03, 7.8125e-03, 7.8125e-03, ..., 3.6065e-319, +0: 3.6065e-319, 3.6066e-319], device='cuda:0', dtype=torch.float64) +0: compute-p: 1.1612800359725952 +0: p.nnz: 72999 +0: pre-loop: 0.19340799748897552 +0: imid: 36499 val: 0.877743 arr[imid]: 1.000000 +0: imid: 18249 val: 0.877743 arr[imid]: 0.452343 +0: imid: 27374 val: 0.877743 arr[imid]: 0.536156 +0: imid: 31936 val: 0.877743 arr[imid]: 0.795653 +0: imid: 34217 val: 0.877743 arr[imid]: 0.974655 +0: imid: 33076 val: 0.877743 arr[imid]: 0.917219 +0: imid: 32506 val: 0.877743 arr[imid]: 0.841953 +0: imid: 32791 val: 0.877743 arr[imid]: 0.890303 +0: imid: 32648 val: 0.877743 arr[imid]: 0.843019 +0: imid: 32719 val: 0.877743 arr[imid]: 0.890276 +0: imid: 32683 val: 0.877743 arr[imid]: 0.890252 +0: imid: 32665 val: 0.877743 arr[imid]: 0.890068 +0: imid: 32656 val: 0.877743 arr[imid]: 0.890068 +0: imid: 32652 val: 0.877743 arr[imid]: 0.843019 +0: imid: 32654 val: 0.877743 arr[imid]: 0.843019 +0: imid: 32655 val: 0.877743 arr[imid]: 0.843019 +0: iter_count: 1 dart_values: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', +0: dtype=torch.float64) +0: iter_count: 1 dart_values_scaled: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', +0: dtype=torch.float64) +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 16], +0: [ 58], +0: [ 87], +0: ..., +0: [36389], +0: [36423], +0: [36479]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([1036, 1]) +0: iter_count: 1 dart_hits_count.sum: 125952 +0: iter_count: 1 dart_hits_map: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', +0: dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[1], +0: [1], +0: [6], +0: ..., +0: [6], +0: [1], +0: [2]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.3693e-12, 2.7386e-12, ..., 1.0000e+00, 1.0000e+00, +0: 1.0000e+00], device='cuda:0', dtype=torch.float64) +0: construct-nextf: 0.2680000066757202 +0: total_time: 389.8416748046875 +0: imid: 18250 val: 0.877743 arr[imid]: 0.452343 +0: imid: 27375 val: 0.877743 arr[imid]: 0.536156 +0: imid: 31938 val: 0.877743 arr[imid]: 0.795653 +0: imid: 34219 val: 0.877743 arr[imid]: 0.974656 +0: imid: 33078 val: 0.877743 arr[imid]: 0.917219 +0: imid: 32508 val: 0.877743 arr[imid]: 0.841953 +0: imid: 32793 val: 0.877743 arr[imid]: 0.890303 +0: imid: 32650 val: 0.877743 arr[imid]: 0.843019 +0: imid: 32721 val: 0.877743 arr[imid]: 0.890276 +0: imid: 32685 val: 0.877743 arr[imid]: 0.890252 +0: imid: 32667 val: 0.877743 arr[imid]: 0.890252 +0: imid: 32658 val: 0.877743 arr[imid]: 0.890068 +0: imid: 32654 val: 0.877743 arr[imid]: 0.843019 +0: imid: 32656 val: 0.877743 arr[imid]: 0.890068 +0: imid: 32655 val: 0.877743 arr[imid]: 0.843019 +0: spgemm-matc-inst total_time: 0.06140799820423126 avg_time 0.06140799820423126 +0: spgemm-bcast-nnz total_time: 123.69152069091797 avg_time 123.69152069091797 +0: spgemm-inst-recv total_time: 4.317759990692139 avg_time 4.317759990692139 +0: spgemm-bcast-data total_time: 0.09142400324344635 avg_time 0.09142400324344635 +0: spgemm-preproc-local total_time: 0.436352014541626 avg_time 0.436352014541626 +0: spgemm-local-spgemm total_time: 170.251708984375 avg_time 170.251708984375 +0: spgemm-chunk-inst total_time: 0.17958399653434753 avg_time 0.17958399653434753 +0: spgemm-chunk-coalesce total_time: 0.49852800369262695 avg_time 0.49852800369262695 +0: spgemm-chunk-add total_time: 0.11667200177907944 avg_time 0.11667200177907944 +0: spgemm-matc-coalesce total_time: 0.011359999887645245 avg_time 0.011359999887645245 +0: spgemm-reduce-nnz total_time: 37.1005744934082 avg_time 37.1005744934082 +0: spgemm-padding total_time: 0.14291200041770935 avg_time 0.14291200041770935 +0: spgemm-allgather total_time: 0.2025279998779297 avg_time 0.2025279998779297 +0: spgemm-preproc-reduce total_time: 0.15545600652694702 avg_time 0.15545600652694702 +0: spgemm-reduce total_time: 1.304703950881958 avg_time 1.304703950881958 +0: spgemm-unpad total_time: 0.26422399282455444 avg_time 0.26422399282455444 +0: prob-rowsum total_time: 0.1865919977426529 avg_time 0.1865919977426529 +0: gen-darts total_time: 0.0759039968252182 avg_time 0.0759039968252182 +0: dart-throw total_time: 1.1985280513763428 avg_time 1.1985280513763428 +0: filter_darts total_time: 2.5079360008239746 avg_time 2.5079360008239746 +0: add-to-frontier total_time: 0.27209600806236267 avg_time 0.27209600806236267 +0: count_samples total_time: 0.13497599959373474 avg_time 0.13497599959373474 +0: select-preproc total_time: 0.1653759926557541 avg_time 0.1653759926557541 +0: select-psoverflow total_time: 0.4081280007958412 avg_time 0.08162560015916824 +0: select-reciprocal total_time: 1.1649919897317886 avg_time 0.23299839794635774 +0: select-instmtx total_time: 0.729408010840416 avg_time 0.14588160216808319 +0: select-invsum total_time: 1.12848000228405 avg_time 0.22569600045681 +0: select-psinv total_time: 0.48076799511909485 avg_time 0.09615359902381897 +0: select-computedarts total_time: 0.30032000318169594 avg_time 0.06006400063633919 +0: select-throwdarts total_time: 0.3787840008735657 avg_time 0.07575680017471313 +0: select-add-to-frontier total_time: 0.9426880180835724 avg_time 0.18853760361671448 +0: select-samplecount total_time: 0.667136013507843 avg_time 0.1334272027015686 +0: select-overflow total_time: 0.4785599932074547 avg_time 0.09571199864149094 +0: select-iter total_time: 7.840160012245178 avg_time 1.5680320024490357 +0: dart-selection total_time: 8.04736042022705 avg_time 8.04736042022705 +0: set-probs total_time: 0.3381119966506958 avg_time 0.3381119966506958 +0: compute-bool total_time: 0.06675200164318085 avg_time 0.06675200164318085 +0: sampling-iters total_time: 13.418304443359375 avg_time 13.418304443359375 +0: select-mtxs total_time: 0.2803199887275696 avg_time 0.2803199887275696 +0: row-select-spgemm total_time: 20.787424087524414 avg_time 20.787424087524414 +0: row-select-expand total_time: 0.10172799974679947 avg_time 0.10172799974679947 +0: col-select-spgemm total_time: 13.91817569732666 avg_time 13.91817569732666 +0: set-sample total_time: 0.12230399996042252 avg_time 0.12230399996042252 +0: iter_count: 1 +0: selection_iter_count: 5 +0: +0: probability-spgemm: 52.053855895996094 +0: before p._values(): tensor([ 7.8125e-03, 7.8125e-03, 7.8125e-03, ..., 2.0000e+00, +0: 2.0000e+00, 5.2635e-315], device='cuda:0', dtype=torch.float64) +0: compute-p: 0.8244479894638062 +0: p.nnz: 36500 +0: pre-loop: 0.16601599752902985 +0: iter_count: 1 dart_values: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', +0: dtype=torch.float64) +0: iter_count: 1 dart_values_scaled: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', +0: dtype=torch.float64) +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 16], +0: [ 58], +0: [ 87], +0: ..., +0: [36389], +0: [36423], +0: [36479]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([1036, 1]) +0: iter_count: 1 dart_hits_count.sum: 125952 +0: iter_count: 1 dart_hits_map: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', +0: dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[1], +0: [1], +0: [6], +0: ..., +0: [6], +0: [1], +0: [2]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.3693e-12, 2.7386e-12, ..., 1.0000e+00, 1.0000e+00, +0: 1.0000e+00], device='cuda:0', dtype=torch.float64) +0: construct-nextf: 0.25760000944137573 +0: total_time: 100.39507293701172 +0: spgemm-matc-inst total_time: 0.039872001856565475 avg_time 0.039872001856565475 +0: spgemm-bcast-nnz total_time: 0.08102399855852127 avg_time 0.08102399855852127 +0: spgemm-inst-recv total_time: 4.348288059234619 avg_time 4.348288059234619 +0: spgemm-bcast-data total_time: 0.061664000153541565 avg_time 0.061664000153541565 +0: spgemm-preproc-local total_time: 0.28252801299095154 avg_time 0.28252801299095154 +0: spgemm-local-spgemm total_time: 44.76665496826172 avg_time 44.76665496826172 +0: spgemm-chunk-inst total_time: 0.14655999839305878 avg_time 0.14655999839305878 +0: spgemm-chunk-coalesce total_time: 0.4310719966888428 avg_time 0.4310719966888428 +0: spgemm-chunk-add total_time: 0.06966400146484375 avg_time 0.06966400146484375 +0: spgemm-matc-coalesce total_time: 0.012160000391304493 avg_time 0.012160000391304493 +0: spgemm-reduce-nnz total_time: 0.09187199920415878 avg_time 0.09187199920415878 +0: spgemm-padding total_time: 0.11664000153541565 avg_time 0.11664000153541565 +0: spgemm-allgather total_time: 0.1693439930677414 avg_time 0.1693439930677414 +0: spgemm-preproc-reduce total_time: 0.15148800611495972 avg_time 0.15148800611495972 +0: spgemm-reduce total_time: 0.6639999747276306 avg_time 0.6639999747276306 +0: spgemm-unpad total_time: 0.22323200106620789 avg_time 0.22323200106620789 +0: prob-rowsum total_time: 0.10636799782514572 avg_time 0.10636799782514572 +0: gen-darts total_time: 0.03984000161290169 avg_time 0.03984000161290169 +0: dart-throw total_time: 1.133247971534729 avg_time 1.133247971534729 +0: filter_darts total_time: 2.332927942276001 avg_time 2.332927942276001 +0: add-to-frontier total_time: 0.1860799938440323 avg_time 0.1860799938440323 +0: count_samples total_time: 0.10454399883747101 avg_time 0.10454399883747101 +0: select-preproc total_time: 0.09487999975681305 avg_time 0.09487999975681305 +0: select-psoverflow total_time: 0.39632000029087067 avg_time 0.07926400005817413 +0: select-reciprocal total_time: 0.8863359987735748 avg_time 0.17726719975471497 +0: select-instmtx total_time: 0.6157120019197464 avg_time 0.12314240038394927 +0: select-invsum total_time: 0.7136320024728775 avg_time 0.1427264004945755 +0: select-psinv total_time: 0.4678719937801361 avg_time 0.09357439875602722 +0: select-computedarts total_time: 0.30399999767541885 avg_time 0.06079999953508377 +0: select-throwdarts total_time: 0.35942399501800537 avg_time 0.07188479900360108 +0: select-add-to-frontier total_time: 0.8053120076656342 avg_time 0.16106240153312684 +0: select-samplecount total_time: 0.48419199883937836 avg_time 0.09683839976787567 +0: select-overflow total_time: 0.46035199612379074 avg_time 0.09207039922475815 +0: select-iter total_time: 6.5983359813690186 avg_time 1.3196671962738038 +0: dart-selection total_time: 6.795360088348389 avg_time 6.795360088348389 +0: set-probs total_time: 0.2848320007324219 avg_time 0.2848320007324219 +0: compute-bool total_time: 0.05846399813890457 avg_time 0.05846399813890457 +0: sampling-iters total_time: 11.527423858642578 avg_time 11.527423858642578 +0: select-mtxs total_time: 0.25360000133514404 avg_time 0.25360000133514404 +0: row-select-spgemm total_time: 20.890464782714844 avg_time 20.890464782714844 +0: row-select-expand total_time: 0.10073599964380264 avg_time 0.10073599964380264 +0: col-select-spgemm total_time: 13.9202880859375 avg_time 13.9202880859375 +0: set-sample total_time: 0.12569600343704224 avg_time 0.12569600343704224 +0: iter_count: 1 +0: selection_iter_count: 5 +0: +0: return_dart_count1.nnzidxs: tensor([ 16, 58, 87, ..., 36389, 36423, 36479], device='cuda:0') +0: return_dart_count1: tensor([1, 1, 6, ..., 6, 1, 2], device='cuda:0', dtype=torch.int32) +0: return_dart_map1: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', +0: dtype=torch.int32) +0: return_dart_val1: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', +0: dtype=torch.float64) +0: return_dart_count2.nnzidxs: tensor([ 16, 58, 87, ..., 36389, 36423, 36479], device='cuda:0') +0: return_dart_count2: tensor([1, 1, 6, ..., 6, 1, 2], device='cuda:0', dtype=torch.int32) +0: return_dart_map2: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', +0: dtype=torch.int32) +0: return_dart_val2: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', +0: dtype=torch.float64) +0: +0: unmatched_darts: tensor([], device='cuda:0', dtype=torch.int64) diff --git a/examples/tmp3.out b/examples/tmp3.out new file mode 100644 index 0000000..d1c2b83 --- /dev/null +++ b/examples/tmp3.out @@ -0,0 +1,473 @@ +0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=512, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) +0: hostname: nid002781 +0: hostname: nid002781 rank: 0 size: 1 +0: start partitioning +0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) +0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) +0: end partitioning +0: coalescing +0: normalizing +0: done normalizing +0: imid: 18166615 val: 8172334.500000 arr[imid]: 2564784128.000000 +0: imid: 9083307 val: 8172334.500000 arr[imid]: 1288451712.000000 +0: imid: 4541653 val: 8172334.500000 arr[imid]: 640568896.000000 +0: imid: 2270826 val: 8172334.500000 arr[imid]: 317143936.000000 +0: imid: 1135412 val: 8172334.500000 arr[imid]: 159784208.000000 +0: imid: 567705 val: 8172334.500000 arr[imid]: 80147880.000000 +0: imid: 283852 val: 8172334.500000 arr[imid]: 41578028.000000 +0: imid: 141925 val: 8172334.500000 arr[imid]: 20473418.000000 +0: imid: 70962 val: 8172334.500000 arr[imid]: 9832408.000000 +0: imid: 35480 val: 8172334.500000 arr[imid]: 5203542.000000 +0: imid: 53221 val: 8172334.500000 arr[imid]: 7887794.000000 +0: imid: 62091 val: 8172334.500000 arr[imid]: 8652051.000000 +0: imid: 57656 val: 8172334.500000 arr[imid]: 8335342.500000 +0: imid: 55438 val: 8172334.500000 arr[imid]: 8064393.500000 +0: imid: 56547 val: 8172334.500000 arr[imid]: 8117769.000000 +0: imid: 57101 val: 8172334.500000 arr[imid]: 8215006.000000 +0: imid: 56824 val: 8172334.500000 arr[imid]: 8170865.500000 +0: imid: 56962 val: 8172334.500000 arr[imid]: 8189993.000000 +0: imid: 56893 val: 8172334.500000 arr[imid]: 8172020.500000 +0: imid: 56927 val: 8172334.500000 arr[imid]: 8188378.500000 +0: imid: 56910 val: 8172334.500000 arr[imid]: 8172235.500000 +0: imid: 56918 val: 8172334.500000 arr[imid]: 8188236.000000 +0: imid: 56914 val: 8172334.500000 arr[imid]: 8172336.000000 +0: imid: 56912 val: 8172334.500000 arr[imid]: 8172327.000000 +0: imid: 56913 val: 8172334.500000 arr[imid]: 8172329.000000 +0: imid: 18166615 val: 3401610.500000 arr[imid]: 0.000000 +0: imid: 27249923 val: 3401610.500000 arr[imid]: 0.000000 +0: imid: 31791577 val: 3401610.500000 arr[imid]: 9985493.000000 +0: imid: 29520750 val: 3401610.500000 arr[imid]: 0.000000 +0: imid: 30656163 val: 3401610.500000 arr[imid]: 0.000000 +0: imid: 31223870 val: 3401610.500000 arr[imid]: 9985493.000000 +0: imid: 30940016 val: 3401610.500000 arr[imid]: 0.000000 +0: imid: 31081943 val: 3401610.500000 arr[imid]: 0.000000 +0: imid: 31152906 val: 3401610.500000 arr[imid]: 6810933.000000 +0: imid: 31117424 val: 3401610.500000 arr[imid]: 2159459.750000 +0: imid: 31135165 val: 3401610.500000 arr[imid]: 4937014.500000 +0: imid: 31126294 val: 3401610.500000 arr[imid]: 3478861.000000 +0: imid: 31121859 val: 3401610.500000 arr[imid]: 2739246.250000 +0: imid: 31124076 val: 3401610.500000 arr[imid]: 3059827.500000 +0: imid: 31125185 val: 3401610.500000 arr[imid]: 3233514.500000 +0: imid: 31125739 val: 3401610.500000 arr[imid]: 3426718.250000 +0: imid: 31125462 val: 3401610.500000 arr[imid]: 3257614.000000 +0: imid: 31125600 val: 3401610.500000 arr[imid]: 3265417.500000 +0: imid: 31125669 val: 3401610.500000 arr[imid]: 3318082.250000 +0: imid: 31125704 val: 3401610.500000 arr[imid]: 3421715.250000 +0: imid: 31125686 val: 3401610.500000 arr[imid]: 3421321.000000 +0: imid: 31125677 val: 3401610.500000 arr[imid]: 3420803.250000 +0: imid: 31125673 val: 3401610.500000 arr[imid]: 3420798.750000 +0: imid: 31125671 val: 3401610.500000 arr[imid]: 3318083.750000 +0: imid: 31125672 val: 3401610.500000 arr[imid]: 3318098.000000 +0: n_darts: 125952 +0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], +0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), +0: values=tensor([2.0568e-07, 2.0568e-07, 2.0568e-07, ..., +0: 4.1311e-06, 4.1311e-06, 4.1311e-06]), +0: device='cuda:0', size=(232965, 232965), nnz=114848857, +0: layout=torch.sparse_coo) +0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 511, 511, 511], +0: [195991, 117560, 100852, ..., 210813, 53334, 131122]]), +0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), +0: device='cuda:0', size=(512, 232965), nnz=131072, layout=torch.sparse_coo) +0: probability-spgemm: 343.3492736816406 +0: before p._values(): tensor([1.2276e+00, 1.8491e-01, 1.8721e+01, ..., 3.6240e-01, 7.9862e+00, +0: 2.4029e+02], device='cuda:0') +0: before p._values()[56912]: 13.246072769165039 p_den[row[56912]]: 6.2924485206604 +0: compute-p: 61.61203384399414 +0: p.nnz: 36333230 +0: pre-loop: 1.4520319700241089 +0: p._values()[56912]: 2.105074405670166 ps_p_values[56912]: 8172327.0 +0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') +0: iter_count: 1 dart_values_scaled: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, +0: 5.1157e+09], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], +0: [ 8], +0: [ 11], +0: ..., +0: [36332465], +0: [36332476], +0: [36332500]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3563590, 1]) +0: iter_count: 1 dart_hits_count.sum: 64487424 +0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[ 1], +0: [ 1], +0: [ 1], +0: ..., +0: [ 9], +0: [ 5], +0: [351]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9509e-01, 2.2447e-01, ..., 5.1202e+09, 5.1202e+09, +0: 5.1202e+09], device='cuda:0') +0: sampled_count: tensor([11024, 9666, 12520, 10538, 12932, 11469, 11153, 10402, 12684, 9842, +0: 12438, 13403, 12459, 10115, 11264, 12120, 13021, 13453, 12467, 11809, +0: 11179, 12947, 11294, 10516, 10519, 11119, 10005, 10003, 9595, 11595, +0: 13284, 11668, 8780, 9898, 12276, 11645, 10163, 10876, 10062, 9497, +0: 10663, 11791, 11570, 12157, 13784, 9359, 11432, 9263, 10236, 10962, +0: 9661, 10819, 10947, 10826, 10793, 10124, 8664, 9424, 10931, 9277, +0: 8969, 10122, 9018, 9295, 10272, 9501, 9365, 10336, 10197, 10276, +0: 8096, 10126, 9304, 9791, 9513, 8806, 8005, 9382, 7457, 8782, +0: 9591, 11029, 9743, 11499, 8958, 8565, 10265, 11149, 10869, 12021, +0: 8819, 9657, 7648, 9854, 10864, 9005, 9849, 9635, 10688, 9500, +0: 10095, 8867, 10907, 11618, 7965, 7262, 9538, 8247, 8774, 7919, +0: 7761, 7956, 7381, 8235, 8723, 9805, 8850, 7419, 8473, 7790, +0: 8735, 9197, 7685, 7017, 8090, 8958, 8467, 9538, 6961, 6832, +0: 9487, 8197, 8245, 7300, 8969, 7602, 7694, 8738, 7283, 7328, +0: 6346, 8730, 8239, 7295, 7483, 8942, 8718, 7367, 7956, 7476, +0: 7024, 7648, 7110, 9080, 8088, 9321, 8696, 8877, 7566, 9046, +0: 8364, 5883, 7471, 8178, 8085, 7091, 8098, 7641, 6458, 7959, +0: 7628, 9135, 9614, 6488, 8197, 7765, 10026, 6966, 7791, 7227, +0: 8740, 7021, 8104, 7890, 7465, 6500, 8742, 6896, 8484, 8183, +0: 8265, 7066, 7550, 6762, 7301, 7924, 7399, 6960, 6635, 8193, +0: 6658, 8946, 9388, 7317, 6856, 5976, 8725, 6323, 8338, 10963, +0: 8423, 6646, 8273, 8477, 7008, 6093, 5807, 5955, 5382, 5805, +0: 7292, 5586, 4903, 5398, 6756, 5422, 5375, 6447, 7429, 8104, +0: 7143, 7807, 6924, 6717, 6088, 5282, 7397, 5732, 5386, 7286, +0: 6578, 6753, 5975, 5473, 5728, 5789, 7533, 6489, 4791, 6969, +0: 6134, 6345, 5086, 5451, 6676, 5303, 5671, 5492, 5682, 5190, +0: 5790, 6892, 4640, 5092, 6053, 4556, 4478, 5817, 6266, 5773, +0: 5562, 5093, 7647, 5904, 5386, 7278, 5140, 4420, 5586, 5943, +0: 5722, 5877, 5422, 5381, 5198, 4538, 5971, 5915, 4892, 5487, +0: 5812, 5632, 5510, 7292, 5447, 4926, 5661, 5477, 6372, 5075, +0: 4935, 4570, 4551, 7161, 7198, 4613, 6546, 5083, 5790, 6444, +0: 6488, 5273, 5317, 6100, 6289, 5295, 5661, 5739, 4919, 6473, +0: 7321, 4467, 5162, 5626, 5163, 5242, 5738, 5566, 5668, 5366, +0: 5975, 6558, 4660, 5576, 4756, 5136, 5945, 5684, 5255, 5117, +0: 6296, 6262, 5105, 7194, 6323, 5272, 5927, 5982, 6396, 6069, +0: 6153, 7047, 6407, 5465, 6407, 5350, 5802, 5449, 5421, 6051, +0: 6494, 6544, 5567, 6552, 5296, 6268, 6327, 5792, 5634, 6637, +0: 5647, 5874, 5290, 4723, 4762, 5722, 5901, 5591, 5204, 6419, +0: 6053, 5675, 4983, 7242, 6072, 5344, 5710, 5781, 6664, 6667, +0: 5141, 5986, 4950, 4950, 4962, 6081, 6078, 4958, 5320, 5592, +0: 6145, 4617, 5423, 6351, 4915, 5882, 5784, 6182, 5397, 4903, +0: 5299, 6336, 5898, 6069, 4632, 5133, 6352, 5503, 5518, 5475, +0: 5775, 5268, 6240, 5027, 5523, 5914, 6114, 6852, 6178, 4491, +0: 4671, 4882, 4856, 4035, 4366, 4375, 4475, 4337, 4132, 4379, +0: 3563, 3770, 4201, 5074, 4138, 4800, 4169, 4887, 5110, 4784, +0: 3523, 4548, 4785, 3886, 4104, 4618, 4601, 4350, 4021, 3934, +0: 3695, 3647, 4122, 3875, 4670, 3810, 4459, 3030, 4273, 3011, +0: 4436, 4376, 4015, 4285, 4554, 4830, 4837, 5278, 4474, 4237, +0: 4983, 4089, 4022, 4586, 4258, 3467, 3603, 3786, 4418, 3655, +0: 3587, 3477, 3542, 5781, 3738, 4358, 3541, 3870, 3507, 3611, +0: 3784, 3311, 4182, 4893, 4322, 3625, 3868, 4655, 4744, 3510, +0: 4002, 3945], device='cuda:0', dtype=torch.int32) +0: p._values()[56912]: 0.0 ps_p_values[56912]: 0.0 +0: iter_count: 2 dart_values: tensor([0.4230, 0.5678, 0.2647, ..., 0.0191, 0.5965, 0.7411], device='cuda:0') +0: iter_count: 2 dart_values_scaled: tensor([4.2305e+06, 5.6784e+06, 2.6466e+06, ..., 5.1102e+09, 5.1160e+09, +0: 5.1174e+09], device='cuda:0') +0: iter_count: 2 dart_hits_count.nnzidxs: tensor([[ 0], +0: [31104767], +0: [31104773], +0: ..., +0: [31172045], +0: [31172053], +0: [31172057]], device='cuda:0') +0: iter_count: 2 dart_hits_count.nnzidxs.size: torch.Size([11976, 1]) +0: iter_count: 2 dart_hits_count.sum: 64487424 +0: iter_count: 2 dart_hits_map: tensor([31129386, 31143585, 31121122, ..., 0, 0, 0], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 2 dart_hits_count: tensor([[64361680], +0: [ 1], +0: [ 1], +0: ..., +0: [ 1], +0: [ 1], +0: [ 1]], device='cuda:0', dtype=torch.int32) +0: iter_count: 2 ps_p_values: tensor([ 0., 0., 0., ..., 9985493., 9985493., 9985493.], +0: device='cuda:0') +0: imid: 18166615 val: 8172334.500000 arr[imid]: 2564784896.000000 +0: imid: 9083307 val: 8172334.500000 arr[imid]: 1288451840.000000 +0: imid: 4541653 val: 8172334.500000 arr[imid]: 640569024.000000 +0: imid: 2270826 val: 8172334.500000 arr[imid]: 317144224.000000 +0: imid: 1135412 val: 8172334.500000 arr[imid]: 159784624.000000 +0: imid: 567705 val: 8172334.500000 arr[imid]: 80148176.000000 +0: imid: 283852 val: 8172334.500000 arr[imid]: 41578188.000000 +0: imid: 141925 val: 8172334.500000 arr[imid]: 20473426.000000 +0: imid: 70962 val: 8172334.500000 arr[imid]: 9832383.000000 +0: imid: 35480 val: 8172334.500000 arr[imid]: 5203529.000000 +0: imid: 53221 val: 8172334.500000 arr[imid]: 7887774.500000 +0: imid: 62091 val: 8172334.500000 arr[imid]: 8652029.000000 +0: imid: 57656 val: 8172334.500000 arr[imid]: 8335322.500000 +0: imid: 55438 val: 8172334.500000 arr[imid]: 8064373.500000 +0: imid: 56547 val: 8172334.500000 arr[imid]: 8117749.500000 +0: imid: 57101 val: 8172334.500000 arr[imid]: 8214986.500000 +0: imid: 56824 val: 8172334.500000 arr[imid]: 8170845.500000 +0: imid: 56962 val: 8172334.500000 arr[imid]: 8189973.500000 +0: imid: 56893 val: 8172334.500000 arr[imid]: 8172001.000000 +0: imid: 56927 val: 8172334.500000 arr[imid]: 8188358.500000 +0: imid: 56910 val: 8172334.500000 arr[imid]: 8172215.500000 +0: imid: 56918 val: 8172334.500000 arr[imid]: 8188216.000000 +0: imid: 56914 val: 8172334.500000 arr[imid]: 8172316.000000 +0: imid: 56916 val: 8172334.500000 arr[imid]: 8172323.500000 +0: imid: 56917 val: 8172334.500000 arr[imid]: 8188216.000000 +0: sampled_count: tensor([ 257, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 11990, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, +0: 256, 256], device='cuda:0', dtype=torch.int32) +0: construct-nextf: 0.7972480058670044 +0: total_time: 3895.928466796875 +0: spgemm-matc-inst total_time: 0.05878400057554245 avg_time 0.05878400057554245 +0: spgemm-bcast-nnz total_time: 122.26163482666016 avg_time 122.26163482666016 +0: spgemm-inst-recv total_time: 3.5645439624786377 avg_time 3.5645439624786377 +0: spgemm-bcast-data total_time: 0.09523200243711472 avg_time 0.09523200243711472 +0: spgemm-preproc-local total_time: 0.3563520014286041 avg_time 0.3563520014286041 +0: spgemm-local-spgemm total_time: 96.65433502197266 avg_time 96.65433502197266 +0: spgemm-chunk-inst total_time: 1.1634880304336548 avg_time 1.1634880304336548 +0: spgemm-chunk-coalesce total_time: 33.33251190185547 avg_time 33.33251190185547 +0: spgemm-chunk-add total_time: 1.1100159883499146 avg_time 1.1100159883499146 +0: spgemm-matc-coalesce total_time: 0.01228800043463707 avg_time 0.01228800043463707 +0: spgemm-reduce-nnz total_time: 36.91321563720703 avg_time 36.91321563720703 +0: spgemm-padding total_time: 1.8552320003509521 avg_time 1.8552320003509521 +0: spgemm-allgather total_time: 2.268160104751587 avg_time 2.268160104751587 +0: spgemm-preproc-reduce total_time: 1.1533440351486206 avg_time 1.1533440351486206 +0: spgemm-reduce total_time: 39.05833435058594 avg_time 39.05833435058594 +0: spgemm-unpad total_time: 3.018752098083496 avg_time 3.018752098083496 +0: prob-rowsum total_time: 1.2051519751548767 avg_time 0.6025759875774384 +0: gen-darts total_time: 0.44383999705314636 avg_time 0.22191999852657318 +0: dart-throw total_time: 3.201184034347534 avg_time 1.600592017173767 +0: filter_darts total_time: 73.3745288848877 avg_time 36.68726444244385 +0: add-to-frontier total_time: 3.6911680698394775 avg_time 1.8455840349197388 +0: count_samples total_time: 61.86393737792969 avg_time 30.931968688964844 +0: select-preproc total_time: 0.2099199965596199 avg_time 0.10495999827980995 +0: select-psoverflow total_time: 1.9985919930040836 avg_time 0.062455999781377614 +0: select-reciprocal total_time: 41.37455999851227 avg_time 1.2929549999535084 +0: select-instmtx total_time: 36.437375664711 avg_time 1.1386679895222187 +0: select-invsum total_time: 1894.122688293457 avg_time 59.19133400917053 +0: select-psinv total_time: 2.2108160331845284 avg_time 0.06908800103701651 +0: select-computedarts total_time: 1.884160004556179 avg_time 0.058880000142380595 +0: select-throwdarts total_time: 10.209279954433441 avg_time 0.31903999857604504 +0: select-add-to-frontier total_time: 53.809184312820435 avg_time 1.6815370097756386 +0: select-samplecount total_time: 991.8761978149414 avg_time 30.99613118171692 +0: select-overflow total_time: 2.233344040811062 avg_time 0.06979200127534568 +0: select-iter total_time: 3043.3382110595703 avg_time 95.10431909561157 +0: dart-selection total_time: 3054.8369750976562 avg_time 1527.4184875488281 +0: set-probs total_time: 3.370400071144104 avg_time 1.685200035572052 +0: compute-bool total_time: 0.10444800183176994 avg_time 0.05222400091588497 +0: sampling-iters total_time: 3206.3609619140625 avg_time 1603.1804809570312 +0: select-mtxs total_time: 0.9584640264511108 avg_time 0.9584640264511108 +0: row-select-spgemm total_time: 198.93040466308594 avg_time 198.93040466308594 +0: row-select-expand total_time: 2.83516788482666 avg_time 2.83516788482666 +0: col-select-spgemm total_time: 79.16544342041016 avg_time 79.16544342041016 +0: set-sample total_time: 0.12639999389648438 avg_time 0.12639999389648438 +0: iter_count: 2 +0: selection_iter_count: 32 +0: +0: probability-spgemm: 184.3711700439453 +0: before p._values(): tensor([1.2276e+00, 1.8491e-01, 1.8721e+01, ..., 3.6240e-01, 7.9862e+00, +0: 2.4029e+02], device='cuda:0') +0: before p._values()[56912]: 13.246072769165039 p_den[row[56912]]: 6.292463779449463 +0: compute-p: 61.37958526611328 +0: p.nnz: 36333230 +0: pre-loop: 1.4438400268554688 +0: p._values()[56912]: 2.105069398880005 ps_p_values[56912]: 8172307.0 +0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') +0: iter_count: 1 dart_values_scaled: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, +0: 5.1157e+09], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], +0: [ 8], +0: [ 11], +0: ..., +0: [36332465], +0: [36332476], +0: [36332500]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3563119, 1]) +0: iter_count: 1 dart_hits_count.sum: 64487424 +0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[ 1], +0: [ 1], +0: [ 1], +0: ..., +0: [ 9], +0: [ 6], +0: [304]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9508e-01, 2.2447e-01, ..., 5.1202e+09, 5.1202e+09, +0: 5.1202e+09], device='cuda:0') +0: sampled_count: tensor([11038, 9667, 12520, 10568, 12768, 11472, 11310, 10395, 12761, 9910, +0: 12425, 13432, 12406, 10202, 11313, 12125, 12968, 13405, 12561, 11745, +0: 11194, 13006, 11227, 10395, 10354, 11110, 10107, 9996, 9598, 11741, +0: 13247, 11627, 8896, 9732, 12241, 11660, 10146, 10866, 10056, 9441, +0: 10619, 11767, 11657, 12170, 13797, 9380, 11389, 9277, 10227, 10911, +0: 9713, 10813, 10883, 10788, 10774, 10147, 8637, 9432, 10923, 9243, +0: 8978, 10170, 9043, 9322, 10205, 9454, 9322, 10353, 10208, 10225, +0: 8144, 10149, 9296, 9792, 9529, 8775, 7987, 9299, 7470, 8734, +0: 9667, 11008, 9733, 11530, 8974, 8481, 10398, 11173, 10894, 12060, +0: 8833, 9604, 7662, 9844, 10962, 9020, 9854, 9650, 10728, 9430, +0: 9986, 8931, 10908, 11575, 7910, 7259, 9472, 8202, 8697, 7909, +0: 7729, 7958, 7362, 8178, 8726, 9826, 8873, 7428, 8491, 7797, +0: 8774, 9180, 7704, 7044, 8096, 8939, 8507, 9518, 6994, 6797, +0: 9494, 8138, 8264, 7249, 8939, 7621, 7669, 8734, 7287, 7316, +0: 6354, 8748, 8245, 7315, 7517, 8989, 8748, 7412, 7965, 7501, +0: 7025, 7581, 7108, 9104, 8063, 9321, 8660, 8851, 7581, 9014, +0: 8373, 5826, 7538, 8144, 8096, 7110, 8038, 7664, 6487, 7958, +0: 7668, 9196, 9613, 6476, 8282, 7659, 10054, 6929, 7765, 7205, +0: 8737, 7037, 8079, 7974, 7485, 6498, 8750, 6987, 8448, 8182, +0: 8264, 7058, 7538, 6759, 7273, 7974, 7404, 7022, 6559, 8117, +0: 6684, 8929, 9311, 7385, 6911, 5942, 8717, 6315, 8271, 10950, +0: 8465, 6670, 8323, 8505, 7032, 6099, 5850, 5951, 5400, 5759, +0: 7351, 5581, 4874, 5359, 6755, 5394, 5361, 6433, 7439, 8102, +0: 7114, 7830, 6920, 6708, 6093, 5274, 7424, 5738, 5379, 7267, +0: 6562, 6738, 5992, 5474, 5726, 5788, 7554, 6488, 4801, 6940, +0: 6136, 6347, 5074, 5442, 6645, 5304, 5738, 5504, 5699, 5204, +0: 5812, 6853, 4638, 5094, 6053, 4599, 4470, 5836, 6226, 5780, +0: 5567, 5083, 7659, 5864, 5411, 7339, 5100, 4415, 5586, 5915, +0: 5691, 5845, 5445, 5372, 5219, 4534, 5999, 5901, 4948, 5483, +0: 5827, 5630, 5559, 7246, 5467, 4884, 5687, 5437, 6376, 5074, +0: 4927, 4594, 4565, 7177, 7211, 4619, 6537, 5087, 5818, 6429, +0: 6468, 5271, 5313, 6125, 6296, 5253, 5661, 5748, 4891, 6444, +0: 7318, 4465, 5150, 5627, 5165, 5245, 5743, 5536, 5678, 5381, +0: 5976, 6579, 4661, 5609, 4757, 5146, 5963, 5703, 5237, 5114, +0: 6305, 6249, 5101, 7171, 6257, 5286, 5902, 6002, 6414, 6107, +0: 6152, 7072, 6410, 5420, 6363, 5397, 5770, 5405, 5389, 6049, +0: 6478, 6551, 5583, 6499, 5315, 6228, 6329, 5793, 5633, 6642, +0: 5634, 5870, 5306, 4701, 4764, 5709, 5925, 5587, 5220, 6395, +0: 6056, 5686, 4966, 7269, 6085, 5358, 5729, 5793, 6671, 6677, +0: 5130, 5999, 4952, 4937, 4914, 6034, 6079, 4927, 5301, 5592, +0: 6145, 4639, 5447, 6324, 4967, 5862, 5804, 6168, 5382, 4937, +0: 5316, 6335, 5866, 6072, 4615, 5088, 6393, 5489, 5479, 5455, +0: 5791, 5293, 6251, 5061, 5485, 5937, 6079, 6906, 6176, 4499, +0: 4670, 4883, 4850, 4041, 4373, 4377, 4486, 4337, 4142, 4381, +0: 3568, 3777, 4208, 5078, 4130, 4806, 4168, 4886, 5110, 4768, +0: 3521, 4537, 4789, 3878, 4083, 4622, 4588, 4356, 4019, 3927, +0: 3691, 3659, 4114, 3882, 4676, 3806, 4460, 3033, 4273, 3014, +0: 4441, 4391, 4001, 4294, 4560, 4828, 4845, 5278, 4482, 4241, +0: 4978, 4077, 4034, 4583, 4253, 3464, 3590, 3784, 4414, 3643, +0: 3593, 3463, 3547, 5781, 3719, 4350, 3535, 3868, 3506, 3610, +0: 3783, 3316, 4171, 4885, 4319, 3630, 3866, 4652, 4751, 3500, +0: 3998, 3950], device='cuda:0', dtype=torch.int32) +0: construct-nextf: 0.785535991191864 +0: total_time: 1534.89404296875 +0: spgemm-matc-inst total_time: 0.0414079986512661 avg_time 0.0414079986512661 +0: spgemm-bcast-nnz total_time: 0.08691199868917465 avg_time 0.08691199868917465 +0: spgemm-inst-recv total_time: 3.5993599891662598 avg_time 3.5993599891662598 +0: spgemm-bcast-data total_time: 0.06860800087451935 avg_time 0.06860800087451935 +0: spgemm-preproc-local total_time: 0.22015999257564545 avg_time 0.22015999257564545 +0: spgemm-local-spgemm total_time: 96.85913848876953 avg_time 96.85913848876953 +0: spgemm-chunk-inst total_time: 1.1541119813919067 avg_time 1.1541119813919067 +0: spgemm-chunk-coalesce total_time: 33.36783981323242 avg_time 33.36783981323242 +0: spgemm-chunk-add total_time: 1.1110399961471558 avg_time 1.1110399961471558 +0: spgemm-matc-coalesce total_time: 0.01228800043463707 avg_time 0.01228800043463707 +0: spgemm-reduce-nnz total_time: 0.08790399879217148 avg_time 0.08790399879217148 +0: spgemm-padding total_time: 1.856160044670105 avg_time 1.856160044670105 +0: spgemm-allgather total_time: 2.243583917617798 avg_time 2.243583917617798 +0: spgemm-preproc-reduce total_time: 1.1478400230407715 avg_time 1.1478400230407715 +0: spgemm-reduce total_time: 39.09904098510742 avg_time 39.09904098510742 +0: spgemm-unpad total_time: 3.023871898651123 avg_time 3.023871898651123 +0: prob-rowsum total_time: 0.5768319964408875 avg_time 0.5768319964408875 +0: gen-darts total_time: 0.20255999267101288 avg_time 0.20255999267101288 +0: dart-throw total_time: 1.5697599649429321 avg_time 1.5697599649429321 +0: filter_darts total_time: 9.821632385253906 avg_time 9.821632385253906 +0: add-to-frontier total_time: 1.8423360586166382 avg_time 1.8423360586166382 +0: count_samples total_time: 30.70307159423828 avg_time 30.70307159423828 +0: select-preproc total_time: 0.07065600156784058 avg_time 0.07065600156784058 +0: select-psoverflow total_time: 0.6276479996740818 avg_time 0.06276479996740818 +0: select-reciprocal total_time: 13.165791749954224 avg_time 1.3165791749954223 +0: select-instmtx total_time: 11.41049611568451 avg_time 1.141049611568451 +0: select-invsum total_time: 591.7541084289551 avg_time 59.17541084289551 +0: select-psinv total_time: 0.6932480111718178 avg_time 0.06932480111718178 +0: select-computedarts total_time: 0.6707199960947037 avg_time 0.06707199960947037 +0: select-throwdarts total_time: 3.707904040813446 avg_time 0.3707904040813446 +0: select-add-to-frontier total_time: 16.77183973789215 avg_time 1.6771839737892151 +0: select-samplecount total_time: 310.00691413879395 avg_time 31.000691413879395 +0: select-overflow total_time: 0.695296011865139 avg_time 0.0695296011865139 +0: select-iter total_time: 951.7676544189453 avg_time 95.17676544189453 +0: dart-selection total_time: 957.3478393554688 avg_time 957.3478393554688 +0: set-probs total_time: 1.7666239738464355 avg_time 1.7666239738464355 +0: compute-bool total_time: 0.04710400104522705 avg_time 0.04710400104522705 +0: sampling-iters total_time: 1005.9694213867188 avg_time 1005.9694213867188 +0: select-mtxs total_time: 0.9820160269737244 avg_time 0.9820160269737244 +0: row-select-spgemm total_time: 197.51199340820312 avg_time 197.51199340820312 +0: row-select-expand total_time: 2.815295934677124 avg_time 2.815295934677124 +0: col-select-spgemm total_time: 79.22278594970703 avg_time 79.22278594970703 +0: set-sample total_time: 0.12403199821710587 avg_time 0.12403199821710587 +0: iter_count: 1 +0: selection_iter_count: 10 +0: +0: return_dart_count1.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], +0: device='cuda:0') +0: return_dart_count1: tensor([ 1, 1, 1, ..., 9, 5, 351], device='cuda:0', +0: dtype=torch.int32) +0: return_dart_map1: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: return_dart_val1: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, +0: 5.1157e+09], device='cuda:0') +0: return_dart_count2.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], +0: device='cuda:0') +0: return_dart_count2: tensor([ 1, 1, 1, ..., 9, 6, 304], device='cuda:0', +0: dtype=torch.int32) +0: return_dart_map2: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: return_dart_val2: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, +0: 5.1157e+09], device='cuda:0') +0: +0: unmatched_darts: tensor([ 12, 43, 49, ..., 64487415, 64487416, 64487417], +0: device='cuda:0') +0: unmatched_darts.size: torch.Size([18767503]) +0: dart_map1[unmatched]: tensor([ 27788, 46136, 40383, ..., 36273593, 36267529, 36273527], +0: device='cuda:0', dtype=torch.int32) +0: dart_map2[unmatched]: tensor([ 27789, 46138, 40388, ..., 36273527, 36267509, 36273468], +0: device='cuda:0', dtype=torch.int32) +0: +0: dart_vals1[unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, +0: 5.1113e+09], device='cuda:0') +0: dart_vals2[unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, +0: 5.1113e+09], device='cuda:0') +0: +0: ps_pvals1[map1unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, +0: 5.1113e+09], device='cuda:0') +0: ps_pvals2[map1unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, +0: 5.1114e+09], device='cuda:0') +0: +0: ps_pvals1[map2unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1113e+09, 5.1103e+09, +0: 5.1113e+09], device='cuda:0') +0: ps_pvals2[map2unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, +0: 5.1113e+09], device='cuda:0') diff --git a/examples/tmp_bsearch_fpoint_error.out b/examples/tmp_bsearch_fpoint_error.out new file mode 100644 index 0000000..7cf3c2d --- /dev/null +++ b/examples/tmp_bsearch_fpoint_error.out @@ -0,0 +1,259 @@ +0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=512, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) +0: hostname: nid003637 +0: hostname: nid003637 rank: 0 size: 1 +0: start partitioning +0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) +0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) +0: end partitioning +0: coalescing +0: normalizing +0: done normalizing +0: imid: 18166615 val: 0.817233 arr[imid]: 256.478424 +0: imid: 9083307 val: 0.817233 arr[imid]: 128.845200 +0: imid: 4541653 val: 0.817233 arr[imid]: 64.056862 +0: imid: 2270826 val: 0.817233 arr[imid]: 31.714394 +0: imid: 1135412 val: 0.817233 arr[imid]: 15.978425 +0: imid: 567705 val: 0.817233 arr[imid]: 8.014808 +0: imid: 283852 val: 0.817233 arr[imid]: 4.157815 +0: imid: 141925 val: 0.817233 arr[imid]: 2.047342 +0: imid: 70962 val: 0.817233 arr[imid]: 0.983242 +0: imid: 35480 val: 0.817233 arr[imid]: 0.520355 +0: imid: 53221 val: 0.817233 arr[imid]: 0.788780 +0: imid: 62091 val: 0.817233 arr[imid]: 0.865206 +0: imid: 57656 val: 0.817233 arr[imid]: 0.833535 +0: imid: 55438 val: 0.817233 arr[imid]: 0.806440 +0: imid: 56547 val: 0.817233 arr[imid]: 0.811778 +0: imid: 57101 val: 0.817233 arr[imid]: 0.821501 +0: imid: 56824 val: 0.817233 arr[imid]: 0.817087 +0: imid: 56962 val: 0.817233 arr[imid]: 0.819000 +0: imid: 56893 val: 0.817233 arr[imid]: 0.817203 +0: imid: 56927 val: 0.817233 arr[imid]: 0.818839 +0: imid: 56910 val: 0.817233 arr[imid]: 0.817225 +0: imid: 56918 val: 0.817233 arr[imid]: 0.818824 +0: imid: 56914 val: 0.817233 arr[imid]: 0.817235 +0: imid: 56912 val: 0.817233 arr[imid]: 0.817234 +0: imid: 56911 val: 0.817233 arr[imid]: 0.817233 +0: imid: 18166615 val: 0.817233 arr[imid]: 256.478424 +0: imid: 9083307 val: 0.817233 arr[imid]: 128.845184 +0: imid: 4541653 val: 0.817233 arr[imid]: 64.056938 +0: imid: 2270826 val: 0.817233 arr[imid]: 31.714397 +0: imid: 1135412 val: 0.817233 arr[imid]: 15.978395 +0: imid: 567705 val: 0.817233 arr[imid]: 8.014786 +0: imid: 283852 val: 0.817233 arr[imid]: 4.157809 +0: imid: 141925 val: 0.817233 arr[imid]: 2.047338 +0: imid: 70962 val: 0.817233 arr[imid]: 0.983241 +0: imid: 35480 val: 0.817233 arr[imid]: 0.520354 +0: imid: 53221 val: 0.817233 arr[imid]: 0.788780 +0: imid: 62091 val: 0.817233 arr[imid]: 0.865205 +0: imid: 57656 val: 0.817233 arr[imid]: 0.833534 +0: imid: 55438 val: 0.817233 arr[imid]: 0.806439 +0: imid: 56547 val: 0.817233 arr[imid]: 0.811777 +0: imid: 57101 val: 0.817233 arr[imid]: 0.821501 +0: imid: 56824 val: 0.817233 arr[imid]: 0.817087 +0: imid: 56962 val: 0.817233 arr[imid]: 0.818999 +0: imid: 56893 val: 0.817233 arr[imid]: 0.817202 +0: imid: 56927 val: 0.817233 arr[imid]: 0.818838 +0: imid: 56910 val: 0.817233 arr[imid]: 0.817224 +0: imid: 56918 val: 0.817233 arr[imid]: 0.818824 +0: imid: 56914 val: 0.817233 arr[imid]: 0.817234 +0: imid: 56912 val: 0.817233 arr[imid]: 0.817233 +0: imid: 56913 val: 0.817233 arr[imid]: 0.817233 +0: n_darts: 125952 +0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], +0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), +0: values=tensor([2.0568e-07, 2.0568e-07, 2.0568e-07, ..., +0: 4.1311e-06, 4.1311e-06, 4.1311e-06]), +0: device='cuda:0', size=(232965, 232965), nnz=114848857, +0: layout=torch.sparse_coo) +0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 511, 511, 511], +0: [195991, 117560, 100852, ..., 210813, 53334, 131122]]), +0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), +0: device='cuda:0', size=(512, 232965), nnz=131072, layout=torch.sparse_coo) +0: probability-spgemm: 352.7101745605469 +0: compute-p: 60.275230407714844 +0: p.nnz: 36333230 +0: pre-loop: 1.4532480239868164 +0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') +0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], +0: [ 8], +0: [ 11], +0: ..., +0: [36332465], +0: [36332476], +0: [36332500]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3814725, 1]) +0: iter_count: 1 dart_hits_count.sum: 64487424 +0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[ 1], +0: [ 1], +0: [ 1], +0: ..., +0: [ 14], +0: [ 3], +0: [274]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9509e-08, 2.2447e-08, ..., 5.1202e+02, 5.1202e+02, +0: 5.1202e+02], device='cuda:0') +0: construct-nextf: 0.7938240170478821 +0: total_time: 1601.1064453125 +0: spgemm-matc-inst total_time: 0.059647999703884125 avg_time 0.059647999703884125 +0: spgemm-bcast-nnz total_time: 132.69155883789062 avg_time 132.69155883789062 +0: spgemm-inst-recv total_time: 3.57206392288208 avg_time 3.57206392288208 +0: spgemm-bcast-data total_time: 0.08755200356245041 avg_time 0.08755200356245041 +0: spgemm-preproc-local total_time: 0.35334399342536926 avg_time 0.35334399342536926 +0: spgemm-local-spgemm total_time: 96.30374145507812 avg_time 96.30374145507812 +0: spgemm-chunk-inst total_time: 1.1650559902191162 avg_time 1.1650559902191162 +0: spgemm-chunk-coalesce total_time: 33.32175827026367 avg_time 33.32175827026367 +0: spgemm-chunk-add total_time: 1.1077120304107666 avg_time 1.1077120304107666 +0: spgemm-matc-coalesce total_time: 0.012768000364303589 avg_time 0.012768000364303589 +0: spgemm-reduce-nnz total_time: 36.258689880371094 avg_time 36.258689880371094 +0: spgemm-padding total_time: 1.8590400218963623 avg_time 1.8590400218963623 +0: spgemm-allgather total_time: 2.2658560276031494 avg_time 2.2658560276031494 +0: spgemm-preproc-reduce total_time: 1.1514240503311157 avg_time 1.1514240503311157 +0: spgemm-reduce total_time: 39.014591217041016 avg_time 39.014591217041016 +0: spgemm-unpad total_time: 3.0386879444122314 avg_time 3.0386879444122314 +0: prob-rowsum total_time: 0.5860480070114136 avg_time 0.5860480070114136 +0: gen-darts total_time: 0.2377600073814392 avg_time 0.2377600073814392 +0: dart-throw total_time: 1.3964799642562866 avg_time 1.3964799642562866 +0: filter_darts total_time: 9.814399719238281 avg_time 9.814399719238281 +0: add-to-frontier total_time: 1.8496639728546143 avg_time 1.8496639728546143 +0: count_samples total_time: 30.425567626953125 avg_time 30.425567626953125 +0: select-preproc total_time: 0.14537599682807922 avg_time 0.14537599682807922 +0: select-psoverflow total_time: 0.5724479965865612 avg_time 0.06360533295406236 +0: select-reciprocal total_time: 11.853407859802246 avg_time 1.317045317755805 +0: select-instmtx total_time: 10.271583795547485 avg_time 1.141287088394165 +0: select-invsum total_time: 531.1580810546875 avg_time 59.01756456163194 +0: select-psinv total_time: 0.6277760118246078 avg_time 0.0697528902027342 +0: select-computedarts total_time: 0.6289280094206333 avg_time 0.06988088993562593 +0: select-throwdarts total_time: 3.477055996656418 avg_time 0.38633955518404645 +0: select-add-to-frontier total_time: 15.113856077194214 avg_time 1.679317341910468 +0: select-samplecount total_time: 277.2289619445801 avg_time 30.80321799384223 +0: select-overflow total_time: 0.6320320069789886 avg_time 0.07022577855322096 +0: select-iter total_time: 853.6518630981445 avg_time 94.85020701090495 +0: dart-selection total_time: 853.9918212890625 avg_time 853.9918212890625 +0: set-probs total_time: 1.7295680046081543 avg_time 1.7295680046081543 +0: compute-bool total_time: 0.055904000997543335 avg_time 0.055904000997543335 +0: sampling-iters total_time: 902.2776489257812 avg_time 902.2776489257812 +0: select-mtxs total_time: 0.9269440174102783 avg_time 0.9269440174102783 +0: row-select-spgemm total_time: 200.27154541015625 avg_time 200.27154541015625 +0: row-select-expand total_time: 2.836479902267456 avg_time 2.836479902267456 +0: col-select-spgemm total_time: 79.11392211914062 avg_time 79.11392211914062 +0: set-sample total_time: 0.131071999669075 avg_time 0.131071999669075 +0: iter_count: 1 +0: selection_iter_count: 9 +0: +0: probability-spgemm: 184.49417114257812 +0: compute-p: 60.3001594543457 +0: p.nnz: 36333230 +0: pre-loop: 1.4463039636611938 +0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') +0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], +0: [ 8], +0: [ 11], +0: ..., +0: [36332465], +0: [36332476], +0: [36332500]], device='cuda:0') +0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3813474, 1]) +0: iter_count: 1 dart_hits_count.sum: 64487424 +0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: iter_count: 1 dart_hits_count: tensor([[ 1], +0: [ 1], +0: [ 1], +0: ..., +0: [ 12], +0: [ 3], +0: [271]], device='cuda:0', dtype=torch.int32) +0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9509e-08, 2.2447e-08, ..., 5.1202e+02, 5.1202e+02, +0: 5.1202e+02], device='cuda:0') +0: construct-nextf: 0.8681280016899109 +0: total_time: 1430.2353515625 +0: spgemm-matc-inst total_time: 0.043327998369932175 avg_time 0.043327998369932175 +0: spgemm-bcast-nnz total_time: 0.0880960002541542 avg_time 0.0880960002541542 +0: spgemm-inst-recv total_time: 3.5996479988098145 avg_time 3.5996479988098145 +0: spgemm-bcast-data total_time: 0.06771200150251389 avg_time 0.06771200150251389 +0: spgemm-preproc-local total_time: 0.2189120054244995 avg_time 0.2189120054244995 +0: spgemm-local-spgemm total_time: 96.97039794921875 avg_time 96.97039794921875 +0: spgemm-chunk-inst total_time: 1.1507200002670288 avg_time 1.1507200002670288 +0: spgemm-chunk-coalesce total_time: 33.33875274658203 avg_time 33.33875274658203 +0: spgemm-chunk-add total_time: 1.1117119789123535 avg_time 1.1117119789123535 +0: spgemm-matc-coalesce total_time: 0.01190400030463934 avg_time 0.01190400030463934 +0: spgemm-reduce-nnz total_time: 0.09206400066614151 avg_time 0.09206400066614151 +0: spgemm-padding total_time: 1.8569920063018799 avg_time 1.8569920063018799 +0: spgemm-allgather total_time: 2.242719888687134 avg_time 2.242719888687134 +0: spgemm-preproc-reduce total_time: 1.1456639766693115 avg_time 1.1456639766693115 +0: spgemm-reduce total_time: 39.106143951416016 avg_time 39.106143951416016 +0: spgemm-unpad total_time: 3.037760019302368 avg_time 3.037760019302368 +0: prob-rowsum total_time: 0.5271360278129578 avg_time 0.5271360278129578 +0: gen-darts total_time: 0.20892800390720367 avg_time 0.20892800390720367 +0: dart-throw total_time: 1.305791974067688 avg_time 1.305791974067688 +0: filter_darts total_time: 9.789440155029297 avg_time 9.789440155029297 +0: add-to-frontier total_time: 1.8480960130691528 avg_time 1.8480960130691528 +0: count_samples total_time: 30.420608520507812 avg_time 30.420608520507812 +0: select-preproc total_time: 0.07011199742555618 avg_time 0.07011199742555618 +0: select-psoverflow total_time: 0.5702720023691654 avg_time 0.06336355581879616 +0: select-reciprocal total_time: 11.712607979774475 avg_time 1.3014008866416082 +0: select-instmtx total_time: 10.291007995605469 avg_time 1.143445332845052 +0: select-invsum total_time: 530.8290901184082 avg_time 58.981010013156464 +0: select-psinv total_time: 0.648159995675087 avg_time 0.07201777729723188 +0: select-computedarts total_time: 0.6293440014123917 avg_time 0.06992711126804352 +0: select-throwdarts total_time: 3.543071985244751 avg_time 0.39367466502719456 +0: select-add-to-frontier total_time: 15.13158392906189 avg_time 1.6812871032290988 +0: select-samplecount total_time: 277.18652725219727 avg_time 30.79850302802192 +0: select-overflow total_time: 0.641087993979454 avg_time 0.07123199933105046 +0: select-iter total_time: 853.2802658081055 avg_time 94.80891842312283 +0: dart-selection total_time: 853.6178588867188 avg_time 853.6178588867188 +0: set-probs total_time: 1.7106560468673706 avg_time 1.7106560468673706 +0: compute-bool total_time: 0.045343998819589615 avg_time 0.045343998819589615 +0: sampling-iters total_time: 901.5848999023438 avg_time 901.5848999023438 +0: select-mtxs total_time: 0.9415680170059204 avg_time 0.9415680170059204 +0: row-select-spgemm total_time: 198.17837524414062 avg_time 198.17837524414062 +0: row-select-expand total_time: 2.817023992538452 avg_time 2.817023992538452 +0: col-select-spgemm total_time: 79.16553497314453 avg_time 79.16553497314453 +0: set-sample total_time: 0.13356800377368927 avg_time 0.13356800377368927 +0: iter_count: 1 +0: selection_iter_count: 9 +0: +0: return_dart_count1.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], +0: device='cuda:0') +0: return_dart_count1: tensor([ 1, 1, 1, ..., 14, 3, 274], device='cuda:0', +0: dtype=torch.int32) +0: return_dart_map1: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: return_dart_val1: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: return_dart_count2.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], +0: device='cuda:0') +0: return_dart_count2: tensor([ 1, 1, 1, ..., 12, 3, 271], device='cuda:0', +0: dtype=torch.int32) +0: return_dart_map2: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], +0: device='cuda:0', dtype=torch.int32) +0: return_dart_val2: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, +0: 5.1157e+02], device='cuda:0') +0: +0: unmatched_darts: tensor([ 92, 93, 106, ..., 64487413, 64487416, 64487417], +0: device='cuda:0') +0: dart_map1[unmatched]: tensor([ 56912, 45262, 50443, ..., 36330476, 36267481, 36273438], +0: device='cuda:0', dtype=torch.int32) +0: dart_map2[unmatched]: tensor([ 56914, 45263, 50444, ..., 36330467, 36267509, 36273436], +0: device='cuda:0', dtype=torch.int32) +0: +0: dart_vals1[unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9604, 511.0260, 511.1345], +0: device='cuda:0') +0: dart_vals2[unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9604, 511.0260, 511.1345], +0: device='cuda:0') +0: +0: ps_pvals1[map1unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9605, 511.0261, 511.1346], +0: device='cuda:0') +0: ps_pvals2[map1unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9606, 511.0259, 511.1347], +0: device='cuda:0') +0: +0: ps_pvals1[map2unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9604, 511.0264, 511.1344], +0: device='cuda:0') +0: ps_pvals2[map2unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9605, 511.0262, 511.1346], +0: device='cuda:0') diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 75c999e..8ec1bdb 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -31,12 +31,30 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} -__device__ int binary_searchf(float *arr, float val, int imin, int imax) { +__device__ int binary_searchf(double *arr, double 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 { @@ -358,8 +376,7 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } -__global__ void ComputeDarts1D(float *dart_values, float *p_rowsum, float *ps_p_rowsum, - int n_darts, int mb_count) { +__global__ void ComputeDarts1D(double *dart_values, int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -367,31 +384,23 @@ __global__ void ComputeDarts1D(float *dart_values, float *p_rowsum, float *ps_p_ int dart_count = n_darts * mb_count; for (int i = id; i < dart_count; i += stride) { int row = i / n_darts; - dart_values[i] *= p_rowsum[row]; - dart_values[i] += ps_p_rowsum[row]; + dart_values[i] += (double)row; } } -void compute_darts1d_gpu(const at::Tensor& dart_values, - const at::Tensor& p_rowsum, - const at::Tensor& ps_p_rowsum, - int n_darts, - int mb_count) { - +void compute_darts1d_gpu(const at::Tensor& dart_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); - ComputeDarts1D<<>>(dart_values.data(), - p_rowsum.data(), - ps_p_rowsum.data(), + ComputeDarts1D<<>>(dart_values.data(), n_darts, mb_count); CHECK_ERROR("dart1d computation error") } -__global__ void ComputeDartsSelect(float *dart_select, float *dart_hits_inv_sum, float *ps_dart_hits_inv_sum, +__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; @@ -416,9 +425,9 @@ void compute_darts_select_gpu(const at::Tensor& dart_select, int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ComputeDartsSelect<<>>(dart_select.data(), - dart_hits_inv_sum.data(), - ps_dart_hits_inv_sum.data(), + ComputeDartsSelect<<>>(dart_select.data(), + dart_hits_inv_sum.data(), + ps_dart_hits_inv_sum.data(), ps_overflow.data(), mb_count, total_overflow); @@ -461,23 +470,32 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, int *h_values, int dart_count, int nnz) { +__global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_values, int *h_map, + 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); + int vtx = -1; + // if (i == 92 || i == 53224) { + if (i == 53224) { + vtx = binary_searchfpf(ps_p_values, dart_values[i], 0, nnz); + } else { + vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz); + } if (vtx < 0 || vtx >= nnz) { printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); - } + } atomicAdd(&h_values[vtx], 1); + h_map[i] = vtx; } } void throw_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& ps_p_values, const at::Tensor& h_values, + const at::Tensor& h_map, int dart_count, int nnz) { @@ -486,15 +504,16 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, 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(), + h_map.data(), dart_count, nnz); CHECK_ERROR("dart throwing error") } -__global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, int *dart_hits_count, +__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; @@ -517,15 +536,15 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ThrowDartsSelect<<>>(dart_select.data(), - ps_dart_hits_inv.data(), + ThrowDartsSelect<<>>(dart_select.data(), + ps_dart_hits_inv.data(), dart_hits_count.data(), total_overflow, nnz); CHECK_ERROR("selection dart throwing error") } -__global__ void Normalize(float *output, float *input, long *index, int 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; @@ -541,7 +560,7 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - Normalize<<>>(output.data(), input.data(), index.data(), len); + Normalize<<>>(output.data(), input.data(), index.data(), len); CHECK_ERROR("normalize error") } From 1a9ba6e7228163474f768ab853ece980172c2922 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 16 Jun 2022 22:17:29 -0700 Subject: [PATCH 072/178] Set padding to -1 instead of 0 for 1.5D allreduce --- cagnet/samplers/ladies.py | 18 +++++++++++++++--- examples/gcn_15d.py | 25 +++++++++++++------------ 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index d4fdc2c..a585244 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -103,7 +103,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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_recv_values.append(torch.cuda.DoubleTensor(matc_nnz).fill_(-1.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.LongTensor(matc_nnz - matc._nnz()).fill_(0.0))) @@ -126,7 +126,8 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-reduce") start_time(start_timer) - nnz_mask = matc._values() != 0 + nnz_mask = matc._values() != -1.0 + print(f"matc.nnz: {matc._nnz()}") matc_nnz_indices = matc._indices()[:, nnz_mask] matc_nnz_values = matc._values()[nnz_mask] stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-unpad") @@ -175,8 +176,16 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") + print(f"before batches: {batches._indices()[1, :]}") + print(f"before adj_matrix: {adj_matrix._indices()}") + print(f"batches.nnz: {batches._nnz()}") + print(f"adj_matrix.nnz: {adj_matrix._nnz()}") p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, row_groups, \ col_groups, timing_dict) + + # p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), + # adj_matrix._indices().long(), adj_matrix._values(), + # mb_count, node_count, node_count, coalesced=True) torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") @@ -189,9 +198,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota size=(mb_count, node_count_total)) print(f"before p._values(): {p._values()}") # print(f"before p._values()[56912]: {p._values()[56912]} p_den[row[56912]]: {p_den[p._indices()[0, 56912]]}") + print(f"p.nnz: {p._nnz()}") + print(f"p.values.size: {p._values().size()}") + print(f"p_num_values.size: {p_num_values.size()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") - print(f"p.nnz: {p._nnz()}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") @@ -387,6 +398,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-bool") underfull_minibatches = (sampled_count < frontier_size).any() + print(f"sampled_count: {sampled_count} frontier_size: {frontier_size} underfull: {underfull_minibatches}") torch.cuda.nvtx.range_pop() timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index c2cf046..3772973 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -135,7 +135,7 @@ def row_normalize(mx): 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 = r_inv._values() + r_inv_values = r_inv._values().double() 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()], @@ -146,7 +146,8 @@ def row_normalize(mx): mx._indices(), mx._values().float(), 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, size=(r_mat_inv.size(0), mx.size(1))) + 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, \ @@ -403,16 +404,16 @@ def main(args): print() unmatched_darts = ((dart_map1 == dart_map2) == 0).nonzero().squeeze() print(f"unmatched_darts: {unmatched_darts}") - # print(f"unmatched_darts.size: {unmatched_darts.size()}") - # print(f"dart_map1[unmatched]: {dart_map1[unmatched_darts]}") - # print(f"dart_map2[unmatched]: {dart_map2[unmatched_darts]}") - # print() - # print(f"dart_vals1[unmatched]: {dart_vals1[unmatched_darts]}") - # print(f"dart_vals2[unmatched]: {dart_vals2[unmatched_darts]}") - # print() - # print(f"ps_pvasl1.size: {ps_pvals1.size()}") - # print(f"ps_pvals1: {ps_pvals1}") - # print(f"ps_pvals2: {ps_pvals2}") + print(f"unmatched_darts.size: {unmatched_darts.size()}") + print(f"dart_map1[unmatched]: {dart_map1[unmatched_darts]}") + print(f"dart_map2[unmatched]: {dart_map2[unmatched_darts]}") + print() + print(f"dart_vals1[unmatched]: {dart_vals1[unmatched_darts]}") + print(f"dart_vals2[unmatched]: {dart_vals2[unmatched_darts]}") + print() + print(f"ps_pvasl1.size: {ps_pvals1.size()}") + print(f"ps_pvals1: {ps_pvals1}") + print(f"ps_pvals2: {ps_pvals2}") # print(f"ps_pvals1[map1unmatched]: {ps_pvals1[dart_map1[unmatched_darts].long()]}") # print(f"ps_pvals2[map1unmatched]: {ps_pvals2[dart_map1[unmatched_darts].long()]}") # print() From 71bfa98900d45405b5ca32ad0a5ca44e9e09c626 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 16 Jun 2022 23:21:05 -0700 Subject: [PATCH 073/178] Fixed long/double mismatch in distributed spgemm --- cagnet/samplers/ladies.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index a585244..29aab10 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -106,12 +106,18 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, matc_recv_values.append(torch.cuda.DoubleTensor(matc_nnz).fill_(-1.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.LongTensor(matc_nnz - matc._nnz()).fill_(0.0))) + matc_send_values = torch.cat((matc._values(), torch.cuda.DoubleTensor(matc_nnz - matc._nnz()).fill_(-1.0))) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-padding") start_time(start_timer) dist.all_gather(matc_recv_indices, matc_send_indices, row_groups[rank_c]) + print(f"before matc_send_values: {matc_send_values}") + print(f"before matc_send_values.size: {matc_send_values.size()}") + print(f"before matc_recv_values[0]: {matc_recv_values[0]}") + print(f"before matc_recv_values[0].size: {matc_recv_values[0].size()}") dist.all_gather(matc_recv_values, matc_send_values, row_groups[rank_c]) + print(f"after matc_recv_values[0]: {matc_recv_values[0]}") + print(f"after matc_recv_values[0].size: {matc_recv_values[0].size()}") stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-allgather") start_time(start_timer) @@ -127,7 +133,6 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, start_time(start_timer) nnz_mask = matc._values() != -1.0 - print(f"matc.nnz: {matc._nnz()}") matc_nnz_indices = matc._indices()[:, nnz_mask] matc_nnz_values = matc._values()[nnz_mask] stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-unpad") @@ -176,10 +181,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - print(f"before batches: {batches._indices()[1, :]}") - print(f"before adj_matrix: {adj_matrix._indices()}") - print(f"batches.nnz: {batches._nnz()}") - print(f"adj_matrix.nnz: {adj_matrix._nnz()}") p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, row_groups, \ col_groups, timing_dict) @@ -188,6 +189,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # mb_count, node_count, node_count, coalesced=True) torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") + print(f"p_num_values: {p_num_values}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") @@ -196,11 +198,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) - print(f"before p._values(): {p._values()}") # print(f"before p._values()[56912]: {p._values()[56912]} p_den[row[56912]]: {p_den[p._indices()[0, 56912]]}") print(f"p.nnz: {p._nnz()}") - print(f"p.values.size: {p._values().size()}") - print(f"p_num_values.size: {p_num_values.size()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") From 52bb6d0ea95c42d709f428cd31e2614a7eeffd91 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 17 Jun 2022 00:26:34 -0700 Subject: [PATCH 074/178] Cleanup --- cagnet/samplers/ladies.py | 23 +-------------- examples/gcn_15d.py | 42 +++++---------------------- sparse-extension/sparse_coo_tensor.cu | 26 +---------------- tests/test_ladies.py | 9 ++++-- 4 files changed, 15 insertions(+), 85 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 29aab10..8a0f6b1 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -111,13 +111,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, start_time(start_timer) dist.all_gather(matc_recv_indices, matc_send_indices, row_groups[rank_c]) - print(f"before matc_send_values: {matc_send_values}") - print(f"before matc_send_values.size: {matc_send_values.size()}") - print(f"before matc_recv_values[0]: {matc_recv_values[0]}") - print(f"before matc_recv_values[0].size: {matc_recv_values[0].size()}") dist.all_gather(matc_recv_values, matc_send_values, row_groups[rank_c]) - print(f"after matc_recv_values[0]: {matc_recv_values[0]}") - print(f"after matc_recv_values[0].size: {matc_recv_values[0].size()}") stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-allgather") start_time(start_timer) @@ -184,12 +178,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, row_groups, \ col_groups, timing_dict) - # p_num_indices, p_num_values = torch_sparse.spspmm(batches._indices(), batches._values(), - # adj_matrix._indices().long(), adj_matrix._values(), - # mb_count, node_count, node_count, coalesced=True) torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") - print(f"p_num_values: {p_num_values}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") @@ -198,7 +188,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) - # print(f"before p._values()[56912]: {p._values()[56912]} p_den[row[56912]]: {p_den[p._indices()[0, 56912]]}") print(f"p.nnz: {p._nnz()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) print(f"compute-p: {stop_time(start_timer, stop_timer)}") @@ -233,7 +222,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-prob-rowsum") ps_p_values = torch.cumsum(p._values().double(), dim=0).roll(1) ps_p_values[0] = 0 - # print(f"p._values()[56912]: {p._values()[56912]} ps_p_values[56912]: {ps_p_values[56912]}") torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) @@ -245,9 +233,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - print(f"iter_count: {iter_count} dart_values: {dart_values}") compute_darts1d_gpu(dart_values, n_darts, mb_count) - print(f"iter_count: {iter_count} dart_values_scaled: {dart_values}") torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) @@ -264,12 +250,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota return_dart_val = dart_values.clone() return_ps_pvals = ps_p_values.clone() - print(f"iter_count: {iter_count} dart_hits_count.nnzidxs: {dart_hits_count.nonzero()}") - print(f"iter_count: {iter_count} dart_hits_count.nnzidxs.size: {dart_hits_count.nonzero().size()}") - print(f"iter_count: {iter_count} dart_hits_count.sum: {torch.sum(dart_hits_count)}") - print(f"iter_count: {iter_count} dart_hits_map: {dart_hits_map}") - print(f"iter_count: {iter_count} dart_hits_count: {dart_hits_count[dart_hits_count.nonzero()]}") - print(f"iter_count: {iter_count} ps_p_values: {ps_p_values}", flush=True) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) @@ -397,7 +377,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-bool") underfull_minibatches = (sampled_count < frontier_size).any() - print(f"sampled_count: {sampled_count} frontier_size: {frontier_size} underfull: {underfull_minibatches}") torch.cuda.nvtx.range_pop() timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) @@ -493,4 +472,4 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") - return batches_select, next_frontier_select, adj_matrices, return_dart_counts, return_dart_map, return_dart_val, return_ps_pvals + return batches_select, next_frontier_select, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 3772973..3b6be48 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -135,19 +135,18 @@ def row_normalize(mx): 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 = r_inv._values().double() + 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().float(), - mx._indices(), mx._values().float(), + mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values().double(), + mx._indices(), mx._values().double(), 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))) + mx = torch.sparse_coo_tensor(indices=mx_indices, values=mx_values, size=(r_mat_inv.size(0), mx.size(1))) return mx def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, \ @@ -326,7 +325,7 @@ def main(args): print("coalescing", flush=True) g_loc = g_loc.coalesce().t() print("normalizing", flush=True) - # g_loc = row_normalize(g_loc) + g_loc = row_normalize(g_loc) print("done normalizing", flush=True) n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) @@ -375,7 +374,7 @@ def main(args): print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing torch.manual_seed(0) - current_frontier, next_frontier, adj_matrices, dart_count1, dart_map1, dart_vals1, ps_pvals1 = \ + current_frontier, next_frontier, adj_matrices = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ @@ -386,39 +385,12 @@ def main(args): torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") torch.manual_seed(0) - current_frontier, next_frontier, adj_matrices_bulk, dart_count2, dart_map2, dart_vals2, ps_pvals2 = \ + 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.replication, rank, size, \ row_groups, col_groups) - print() - print(f"return_dart_count1.nnzidxs: {dart_count1.nonzero().squeeze()}") - print(f"return_dart_count1: {dart_count1[dart_count1.nonzero().squeeze()]}") - print(f"return_dart_map1: {dart_map1}") - print(f"return_dart_val1: {dart_vals1}") - print(f"return_dart_count2.nnzidxs: {dart_count2.nonzero().squeeze()}") - print(f"return_dart_count2: {dart_count2[dart_count2.nonzero().squeeze()]}") - print(f"return_dart_map2: {dart_map2}") - print(f"return_dart_val2: {dart_vals2}") - print() - unmatched_darts = ((dart_map1 == dart_map2) == 0).nonzero().squeeze() - print(f"unmatched_darts: {unmatched_darts}") - print(f"unmatched_darts.size: {unmatched_darts.size()}") - print(f"dart_map1[unmatched]: {dart_map1[unmatched_darts]}") - print(f"dart_map2[unmatched]: {dart_map2[unmatched_darts]}") - print() - print(f"dart_vals1[unmatched]: {dart_vals1[unmatched_darts]}") - print(f"dart_vals2[unmatched]: {dart_vals2[unmatched_darts]}") - print() - print(f"ps_pvasl1.size: {ps_pvals1.size()}") - print(f"ps_pvals1: {ps_pvals1}") - print(f"ps_pvals2: {ps_pvals2}") - # print(f"ps_pvals1[map1unmatched]: {ps_pvals1[dart_map1[unmatched_darts].long()]}") - # print(f"ps_pvals2[map1unmatched]: {ps_pvals2[dart_map1[unmatched_darts].long()]}") - # print() - # print(f"ps_pvals1[map2unmatched]: {ps_pvals1[dart_map2[unmatched_darts].long()]}") - # print(f"ps_pvals2[map2unmatched]: {ps_pvals2[dart_map2[unmatched_darts].long()]}") torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 8ec1bdb..1f1ad4e 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -48,24 +48,6 @@ __device__ int binary_searchf(double *arr, double val, int imin, int imax) { 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; @@ -477,13 +459,7 @@ __global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_va int stride = blockDim.x * gridDim.x; for (int i = id; i < dart_count; i += stride) { - int vtx = -1; - // if (i == 92 || i == 53224) { - if (i == 53224) { - vtx = binary_searchfpf(ps_p_values, dart_values[i], 0, nnz); - } else { - vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz); - } + int vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz); if (vtx < 0 || vtx >= nnz) { printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); } diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 89a3f43..828a4cc 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -1,11 +1,14 @@ import argparse import examples.gcn_15d -import LADIES.pytorch_ladies +# import LADIES.pytorch_ladies +import math import numpy as np import torch import torch.distributed as dist +eps = 1e-4 + def main(args): """ # Testing sampled neighbors, probability distribution, and normalized Laplacian @@ -120,7 +123,7 @@ def main(args): if v_vtx in neighbors: v_vtx_idx = (neighbors == v_vtx).nonzero().item() - if adj_matrix._values()[nnz_count] != neighbors_nnz[v_vtx_idx]: + 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 @@ -164,7 +167,7 @@ def main(args): 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, + parser.add_argument('--n-darts', type=int, default=-1, help='Number of darts to throw per minibatch in LADIES sampling') args = parser.parse_args() print(args) From f87628a573c44be2e8812fc679c687842807a5a7 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sun, 19 Jun 2022 00:21:50 -0700 Subject: [PATCH 075/178] Cleanup and use 0 instead of -1 as padding --- cagnet/samplers/ladies.py | 18 ++++++++++++------ examples/gcn_15d.py | 28 +++++++++++++++------------- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 8a0f6b1..5f8968f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -4,9 +4,10 @@ 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, shift_rowselect_gpu, shift_colselect_gpu +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 def start_time(timer): timer.record() @@ -20,7 +21,9 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if timing_dict is not None: timing_dict[range_name].append(stop_time(start_timer, stop_timer)) -def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, timing_dict=None): +def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ + timing_dict=None): + start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -103,10 +106,10 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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_(-1.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_(-1.0))) + 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, "spgemm-padding") start_time(start_timer) @@ -180,6 +183,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") + print(f"p_num_values.size: {p_num_values.size()}") + print(f"p_num_values: {p_num_values}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") @@ -214,6 +219,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota return_ps_pvals = None while underfull_minibatches: + print(f"iter_count: {iter_count}") iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 3b6be48..11c73a0 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -376,10 +376,10 @@ def main(args): torch.manual_seed(0) current_frontier, next_frontier, adj_matrices = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - args.replication, rank, size, \ - row_groups, col_groups) + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, rank, size, \ + row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() @@ -387,10 +387,10 @@ def main(args): torch.manual_seed(0) 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.replication, rank, size, \ - row_groups, col_groups) + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, rank, size, \ + row_groups, col_groups) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() @@ -411,19 +411,21 @@ def main(args): adj_matrix_sample_indices[0,:] -= row_select_min adj_matrix_sample_values = sampled_values[sample_select_mask] - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, adj_matrix_sample_values, \ - size=(args.batch_size, args.samp_num + args.batch_size)) + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) adj_matrices[j][i] = adj_matrix_sample # print(f"sample: {adj_matrices}") 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)) + torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0), + size=(node_count, node_count)) adj_matrix = adj_matrix.coalesce().cuda() adj_matrix = row_normalize(adj_matrix) adj_matrix = torch.pow(adj_matrix, 2) - return current_frontier, next_frontier, adj_matrices, adj_matrix # return here while testing sampling code + # return here while testing sampling code + return current_frontier, next_frontier, adj_matrices, adj_matrix # create GCN model torch.manual_seed(0) From a4d745d895ceb869d6020791b9b8017e32ac5673 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 24 Jun 2022 22:11:19 -0700 Subject: [PATCH 076/178] Additional cleanup --- cagnet/samplers/ladies.py | 31 ++- examples/tmp.out | 263 -------------------------- sparse-extension/sparse_coo_tensor.cu | 48 +++++ 3 files changed, 61 insertions(+), 281 deletions(-) delete mode 100644 examples/tmp.out diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 5f8968f..199a141 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -7,7 +7,8 @@ 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 + shift_rowselect_gpu, shift_colselect_gpu, \ + scatterd_add_gpu, scatteri_add_gpu def start_time(timer): timer.record() @@ -183,13 +184,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") - print(f"p_num_values.size: {p_num_values.size()}") - print(f"p_num_values: {p_num_values}") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) - p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) + # p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) + scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) @@ -213,13 +213,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True - return_dart_counts = None - return_dart_map = None - return_dart_val = None - return_ps_pvals = None while underfull_minibatches: - print(f"iter_count: {iter_count}") iter_count += 1 start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -250,12 +245,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota dart_hits_map = torch.cuda.IntTensor(n_darts * mb_count) throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, dart_hits_map, \ n_darts * mb_count, p._nnz()) - if iter_count == 1: - return_dart_counts = dart_hits_count.clone() - return_dart_map = dart_hits_map.clone() - return_dart_val = dart_values.clone() - return_ps_pvals = ps_p_values.clone() - torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) @@ -272,7 +261,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-count-samples") # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() sampled_count = sampled_count.fill_(0) - sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0, :], next_frontier_values) + # sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0, :], next_frontier_values) + scatteri_add_gpu(sampled_count, next_frontier_tmp._indices()[0, :], next_frontier_values, \ + next_frontier_values.size(0)) torch.cuda.nvtx.range_pop() timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) @@ -313,7 +304,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-select-invsum") # dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) - dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) + # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) + scatterd_add_gpu(dart_hits_inv_sum, next_frontier._indices()[0, :], dart_hits_inv, \ + dart_hits_inv.size(0)) timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) @@ -350,7 +343,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-select-samplecount") # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) - sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values()) + # sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values()) + scatteri_add_gpu(sampled_count, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values(), \ + next_frontier_tmp._nnz()) timing_dict["select-samplecount"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) diff --git a/examples/tmp.out b/examples/tmp.out deleted file mode 100644 index 9698812..0000000 --- a/examples/tmp.out +++ /dev/null @@ -1,263 +0,0 @@ -0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=512, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) -0: hostname: nid002964 -0: hostname: nid002964 rank: 0 size: 1 -0: start partitioning -0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) -0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) -0: end partitioning -0: coalescing -0: normalizing -0: done normalizing -0: imid: 18166615 val: 0.817233 arr[imid]: 256.417358 -0: imid: 9083307 val: 0.817233 arr[imid]: 128.898560 -0: imid: 4541653 val: 0.817233 arr[imid]: 64.081360 -0: imid: 2270826 val: 0.817233 arr[imid]: 31.746990 -0: imid: 1135412 val: 0.817233 arr[imid]: 15.978256 -0: imid: 567705 val: 0.817233 arr[imid]: 8.033811 -0: imid: 283852 val: 0.817233 arr[imid]: 4.154286 -0: imid: 141925 val: 0.817233 arr[imid]: 2.080449 -0: imid: 70962 val: 0.817233 arr[imid]: 0.972775 -0: imid: 35480 val: 0.817233 arr[imid]: 0.488077 -0: imid: 53221 val: 0.817233 arr[imid]: 0.730208 -0: imid: 62091 val: 0.817233 arr[imid]: 0.852156 -0: imid: 57656 val: 0.817233 arr[imid]: 0.791403 -0: imid: 59873 val: 0.817233 arr[imid]: 0.820989 -0: imid: 58764 val: 0.817233 arr[imid]: 0.806108 -0: imid: 59318 val: 0.817233 arr[imid]: 0.813626 -0: imid: 59595 val: 0.817233 arr[imid]: 0.817638 -0: imid: 59456 val: 0.817233 arr[imid]: 0.815839 -0: imid: 59525 val: 0.817233 arr[imid]: 0.816724 -0: imid: 59560 val: 0.817233 arr[imid]: 0.817125 -0: imid: 59577 val: 0.817233 arr[imid]: 0.817343 -0: imid: 59568 val: 0.817233 arr[imid]: 0.817230 -0: imid: 59572 val: 0.817233 arr[imid]: 0.817272 -0: imid: 59570 val: 0.817233 arr[imid]: 0.817251 -0: imid: 59569 val: 0.817233 arr[imid]: 0.817237 -0: imid: 18166615 val: 0.817233 arr[imid]: 256.417145 -0: imid: 9083307 val: 0.817233 arr[imid]: 128.898499 -0: imid: 4541653 val: 0.817233 arr[imid]: 64.081367 -0: imid: 2270826 val: 0.817233 arr[imid]: 31.746983 -0: imid: 1135412 val: 0.817233 arr[imid]: 15.978255 -0: imid: 567705 val: 0.817233 arr[imid]: 8.033811 -0: imid: 283852 val: 0.817233 arr[imid]: 4.154286 -0: imid: 141925 val: 0.817233 arr[imid]: 2.080449 -0: imid: 70962 val: 0.817233 arr[imid]: 0.972775 -0: imid: 35480 val: 0.817233 arr[imid]: 0.488077 -0: imid: 53221 val: 0.817233 arr[imid]: 0.730208 -0: imid: 62091 val: 0.817233 arr[imid]: 0.852156 -0: imid: 57656 val: 0.817233 arr[imid]: 0.791403 -0: imid: 59873 val: 0.817233 arr[imid]: 0.820989 -0: imid: 58764 val: 0.817233 arr[imid]: 0.806108 -0: imid: 59318 val: 0.817233 arr[imid]: 0.813626 -0: imid: 59595 val: 0.817233 arr[imid]: 0.817638 -0: imid: 59456 val: 0.817233 arr[imid]: 0.815839 -0: imid: 59525 val: 0.817233 arr[imid]: 0.816724 -0: imid: 59560 val: 0.817233 arr[imid]: 0.817125 -0: imid: 59577 val: 0.817233 arr[imid]: 0.817343 -0: imid: 59568 val: 0.817233 arr[imid]: 0.817230 -0: imid: 59572 val: 0.817233 arr[imid]: 0.817272 -0: imid: 59570 val: 0.817233 arr[imid]: 0.817251 -0: imid: 59569 val: 0.817233 arr[imid]: 0.817237 -0: n_darts: 125952 -0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], -0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), -0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), -0: device='cuda:0', size=(232965, 232965), nnz=114848857, -0: layout=torch.sparse_coo) -0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 511, 511, 511], -0: [195991, 117560, 100852, ..., 210813, 53334, 131122]]), -0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), -0: device='cuda:0', size=(512, 232965), nnz=131072, layout=torch.sparse_coo) -0: probability-spgemm: 473.3919677734375 -0: before p._values(): tensor([3., 1., 3., ..., 1., 1., 1.], device='cuda:0') -0: before p._values()[56912]: 2.0 p_den[row[56912]]: 142332.0 -0: compute-p: 61.10249710083008 -0: p.nnz: 36333230 -0: pre-loop: 1.451200008392334 -0: p._values()[56912]: 1.4051654034119565e-05 ps_p_values[56912]: 0.7808363437652588 -0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') -0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 1], -0: [ 3], -0: [ 4], -0: ..., -0: [36333219], -0: [36333221], -0: [36333226]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([18000704, 1]) -0: iter_count: 1 dart_hits_count.sum: 64487424 -0: iter_count: 1 dart_hits_map: tensor([ 28897, 37545, 1807, ..., 36289331, 36332750, 36304529], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[3], -0: [1], -0: [1], -0: ..., -0: [2], -0: [3], -0: [2]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 2.1077e-05, 2.8103e-05, ..., 5.1200e+02, 5.1200e+02, -0: 5.1200e+02], device='cuda:0') -0: construct-nextf: 0.834559977054596 -0: total_time: 1820.016845703125 -0: spgemm-matc-inst total_time: 0.060736000537872314 avg_time 0.060736000537872314 -0: spgemm-bcast-nnz total_time: 123.55779266357422 avg_time 123.55779266357422 -0: spgemm-inst-recv total_time: 3.5558080673217773 avg_time 3.5558080673217773 -0: spgemm-bcast-data total_time: 0.09337600320577621 avg_time 0.09337600320577621 -0: spgemm-preproc-local total_time: 0.5993279814720154 avg_time 0.5993279814720154 -0: spgemm-local-spgemm total_time: 223.27728271484375 avg_time 223.27728271484375 -0: spgemm-chunk-inst total_time: 1.167680025100708 avg_time 1.167680025100708 -0: spgemm-chunk-coalesce total_time: 33.264705657958984 avg_time 33.264705657958984 -0: spgemm-chunk-add total_time: 1.1167999505996704 avg_time 1.1167999505996704 -0: spgemm-matc-coalesce total_time: 0.01235199999064207 avg_time 0.01235199999064207 -0: spgemm-reduce-nnz total_time: 37.626495361328125 avg_time 37.626495361328125 -0: spgemm-padding total_time: 1.8718719482421875 avg_time 1.8718719482421875 -0: spgemm-allgather total_time: 2.272831916809082 avg_time 2.272831916809082 -0: spgemm-preproc-reduce total_time: 1.1596800088882446 avg_time 1.1596800088882446 -0: spgemm-reduce total_time: 40.225215911865234 avg_time 40.225215911865234 -0: spgemm-unpad total_time: 3.052704095840454 avg_time 3.052704095840454 -0: prob-rowsum total_time: 0.612608015537262 avg_time 0.612608015537262 -0: gen-darts total_time: 0.23974399268627167 avg_time 0.23974399268627167 -0: dart-throw total_time: 1.302880048751831 avg_time 1.302880048751831 -0: filter_darts total_time: 9.320735931396484 avg_time 9.320735931396484 -0: add-to-frontier total_time: 1.856287956237793 avg_time 1.856287956237793 -0: count_samples total_time: 30.434911727905273 avg_time 30.434911727905273 -0: select-preproc total_time: 0.15148800611495972 avg_time 0.15148800611495972 -0: select-psoverflow total_time: 0.6426879912614822 avg_time 0.06426879912614822 -0: select-reciprocal total_time: 13.074944019317627 avg_time 1.3074944019317627 -0: select-instmtx total_time: 11.400192022323608 avg_time 1.1400192022323608 -0: select-invsum total_time: 589.2849617004395 avg_time 58.92849617004394 -0: select-psinv total_time: 0.7217599973082542 avg_time 0.07217599973082542 -0: select-computedarts total_time: 1.292479995638132 avg_time 0.12924799956381322 -0: select-throwdarts total_time: 5.817088067531586 avg_time 0.5817088067531586 -0: select-add-to-frontier total_time: 16.777632117271423 avg_time 1.6777632117271424 -0: select-samplecount total_time: 307.49686431884766 avg_time 30.749686431884765 -0: select-overflow total_time: 0.7064639925956726 avg_time 0.07064639925956726 -0: select-iter total_time: 949.5877990722656 avg_time 94.95877990722656 -0: dart-selection total_time: 949.9708251953125 avg_time 949.9708251953125 -0: set-probs total_time: 1.7460800409317017 avg_time 1.7460800409317017 -0: compute-bool total_time: 0.056543998420238495 avg_time 0.056543998420238495 -0: sampling-iters total_time: 997.7376098632812 avg_time 997.7376098632812 -0: select-mtxs total_time: 1.0613759756088257 avg_time 1.0613759756088257 -0: row-select-spgemm total_time: 201.4571533203125 avg_time 201.4571533203125 -0: row-select-expand total_time: 2.8359999656677246 avg_time 2.8359999656677246 -0: col-select-spgemm total_time: 79.69513702392578 avg_time 79.69513702392578 -0: set-sample total_time: 0.1319040060043335 avg_time 0.1319040060043335 -0: iter_count: 1 -0: selection_iter_count: 10 -0: -0: probability-spgemm: 184.28134155273438 -0: before p._values(): tensor([3., 1., 3., ..., 1., 1., 1.], device='cuda:0') -0: before p._values()[56912]: 2.0 p_den[row[56912]]: 142332.0 -0: compute-p: 60.85647964477539 -0: p.nnz: 36333230 -0: pre-loop: 1.4413119554519653 -0: p._values()[56912]: 1.4051654034119565e-05 ps_p_values[56912]: 0.7808363437652588 -0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') -0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 0], -0: [ 1], -0: [ 3], -0: ..., -0: [36333219], -0: [36333221], -0: [36333224]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([17999426, 1]) -0: iter_count: 1 dart_hits_count.sum: 64487424 -0: iter_count: 1 dart_hits_map: tensor([ 28897, 37545, 1807, ..., 36289345, 36332774, 36304546], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[36], -0: [ 3], -0: [ 1], -0: ..., -0: [ 4], -0: [ 3], -0: [ 2]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 2.1077e-05, 2.8103e-05, ..., 5.1200e+02, 5.1200e+02, -0: 5.1200e+02], device='cuda:0') -0: construct-nextf: 0.7904639840126038 -0: total_time: 1524.7774658203125 -0: spgemm-matc-inst total_time: 0.04438399896025658 avg_time 0.04438399896025658 -0: spgemm-bcast-nnz total_time: 0.08675199747085571 avg_time 0.08675199747085571 -0: spgemm-inst-recv total_time: 3.5789759159088135 avg_time 3.5789759159088135 -0: spgemm-bcast-data total_time: 0.06678400188684464 avg_time 0.06678400188684464 -0: spgemm-preproc-local total_time: 0.22096000611782074 avg_time 0.22096000611782074 -0: spgemm-local-spgemm total_time: 96.86115264892578 avg_time 96.86115264892578 -0: spgemm-chunk-inst total_time: 1.154528021812439 avg_time 1.154528021812439 -0: spgemm-chunk-coalesce total_time: 33.27449417114258 avg_time 33.27449417114258 -0: spgemm-chunk-add total_time: 1.1093759536743164 avg_time 1.1093759536743164 -0: spgemm-matc-coalesce total_time: 0.012256000190973282 avg_time 0.012256000190973282 -0: spgemm-reduce-nnz total_time: 0.09273599833250046 avg_time 0.09273599833250046 -0: spgemm-padding total_time: 1.8496320247650146 avg_time 1.8496320247650146 -0: spgemm-allgather total_time: 2.2429440021514893 avg_time 2.2429440021514893 -0: spgemm-preproc-reduce total_time: 1.153056025505066 avg_time 1.153056025505066 -0: spgemm-reduce total_time: 39.10028839111328 avg_time 39.10028839111328 -0: spgemm-unpad total_time: 3.0185279846191406 avg_time 3.0185279846191406 -0: prob-rowsum total_time: 0.5772160291671753 avg_time 0.5772160291671753 -0: gen-darts total_time: 0.21264000236988068 avg_time 0.21264000236988068 -0: dart-throw total_time: 1.2690880298614502 avg_time 1.2690880298614502 -0: filter_darts total_time: 9.313535690307617 avg_time 9.313535690307617 -0: add-to-frontier total_time: 1.840448021888733 avg_time 1.840448021888733 -0: count_samples total_time: 30.42483139038086 avg_time 30.42483139038086 -0: select-preproc total_time: 0.07369600236415863 avg_time 0.07369600236415863 -0: select-psoverflow total_time: 0.6427519991993904 avg_time 0.06427519991993905 -0: select-reciprocal total_time: 13.018144130706787 avg_time 1.3018144130706788 -0: select-instmtx total_time: 11.403391718864441 avg_time 1.1403391718864442 -0: select-invsum total_time: 589.0520629882812 avg_time 58.905206298828126 -0: select-psinv total_time: 0.7290239855647087 avg_time 0.07290239855647088 -0: select-computedarts total_time: 1.2980800233781338 avg_time 0.12980800233781337 -0: select-throwdarts total_time: 5.81113600730896 avg_time 0.581113600730896 -0: select-add-to-frontier total_time: 16.790783882141113 avg_time 1.6790783882141114 -0: select-samplecount total_time: 307.7421760559082 avg_time 30.77421760559082 -0: select-overflow total_time: 0.7125119939446449 avg_time 0.07125119939446449 -0: select-iter total_time: 949.5481796264648 avg_time 94.95481796264649 -0: dart-selection total_time: 949.9288330078125 avg_time 949.9288330078125 -0: set-probs total_time: 1.6976959705352783 avg_time 1.6976959705352783 -0: compute-bool total_time: 0.045791998505592346 avg_time 0.045791998505592346 -0: sampling-iters total_time: 997.4190063476562 avg_time 997.4190063476562 -0: select-mtxs total_time: 0.9089279770851135 avg_time 0.9089279770851135 -0: row-select-spgemm total_time: 196.4358367919922 avg_time 196.4358367919922 -0: row-select-expand total_time: 2.838144063949585 avg_time 2.838144063949585 -0: col-select-spgemm total_time: 79.38406372070312 avg_time 79.38406372070312 -0: set-sample total_time: 0.12310399860143661 avg_time 0.12310399860143661 -0: iter_count: 1 -0: selection_iter_count: 10 -0: -0: return_dart_count1.nnzidxs: tensor([ 1, 3, 4, ..., 36333219, 36333221, 36333226], -0: device='cuda:0') -0: return_dart_count1: tensor([3, 1, 1, ..., 2, 3, 2], device='cuda:0', dtype=torch.int32) -0: return_dart_map1: tensor([ 28897, 37545, 1807, ..., 36289331, 36332750, 36304529], -0: device='cuda:0', dtype=torch.int32) -0: return_dart_val1: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: return_dart_count2.nnzidxs: tensor([ 0, 1, 3, ..., 36333219, 36333221, 36333224], -0: device='cuda:0') -0: return_dart_count2: tensor([36, 3, 1, ..., 4, 3, 2], device='cuda:0', dtype=torch.int32) -0: return_dart_map2: tensor([ 28897, 37545, 1807, ..., 36289345, 36332774, 36304546], -0: device='cuda:0', dtype=torch.int32) -0: return_dart_val2: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: -0: unmatched_darts: tensor([ 6396, 6582, 7501, ..., 64487421, 64487422, 64487423], -0: device='cuda:0') -0: unmatched_darts.size: torch.Size([53266787]) -0: dart_map1[unmatched]: tensor([ 68599, 58618, 35396, ..., 36289331, 36332750, 36304529], -0: device='cuda:0', dtype=torch.int32) -0: dart_map2[unmatched]: tensor([ 68598, 58619, 35397, ..., 36289345, 36332774, 36304546], -0: device='cuda:0', dtype=torch.int32) -0: -0: dart_vals1[unmatched]: tensor([9.4084e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: dart_vals2[unmatched]: tensor([9.4084e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: -0: ps_pvals1[map1unmatched]: tensor([9.4086e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: ps_pvals2[map1unmatched]: tensor([9.4086e-01, 8.0435e-01, 4.8709e-01, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: -0: ps_pvals1[map2unmatched]: tensor([9.4084e-01, 8.0436e-01, 4.8711e-01, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: ps_pvals2[map2unmatched]: tensor([9.4084e-01, 8.0436e-01, 4.8711e-01, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 1f1ad4e..226dc8c 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -597,6 +597,52 @@ void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, i 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") +} + 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"); @@ -610,4 +656,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("normalize_gpu", &normalize_gpu, "Normalize values in an array based on a second array"); 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"); } From e44952f58d344d9c8c9aea0d59baec92ad30b243 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 28 Jun 2022 10:56:06 -0700 Subject: [PATCH 077/178] Additional double->float conversion --- cagnet/samplers/ladies.py | 28 ++++++++--------- examples/gcn_15d.py | 10 +++--- sparse-extension/sparse_coo_tensor.cu | 44 +++++++++++++-------------- 3 files changed, 41 insertions(+), 41 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 199a141..2de72a0 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -8,7 +8,7 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu + scatterf_add_gpu, scatteri_add_gpu def start_time(timer): timer.record() @@ -51,7 +51,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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()) + matb_recv_values = torch.cuda.FloatTensor(matb_recv_nnz.item()) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-inst-recv") start_time(start_timer) @@ -107,10 +107,10 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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_recv_values.append(torch.cuda.FloatTensor(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))) + matc_send_values = torch.cat((matc._values(), torch.cuda.FloatTensor(matc_nnz - matc._nnz()).fill_(0.0))) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-padding") start_time(start_timer) @@ -187,9 +187,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) + p_den = torch.cuda.FloatTensor(mb_count).fill_(0) # p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) - scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + scatterf_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) @@ -221,14 +221,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-prob-rowsum") - ps_p_values = torch.cumsum(p._values().double(), dim=0).roll(1) + ps_p_values = torch.cumsum(p._values().float(), dim=0).roll(1) ps_p_values[0] = 0 torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") - dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() + dart_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) @@ -289,7 +289,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-reciprocal") - dart_hits_inv = dart_hits_count.reciprocal().double() + dart_hits_inv = dart_hits_count.reciprocal().float() dart_hits_inv[dart_hits_inv == float("inf")] = 0 timing_dict["select-reciprocal"].append(stop_time(start_timer, stop_timer)) @@ -303,9 +303,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-invsum") # dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() - dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) + dart_hits_inv_sum = torch.cuda.FloatTensor(mb_count).fill_(0) # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) - scatterd_add_gpu(dart_hits_inv_sum, next_frontier._indices()[0, :], dart_hits_inv, \ + scatterf_add_gpu(dart_hits_inv_sum, next_frontier._indices()[0, :], dart_hits_inv, \ dart_hits_inv.size(0)) timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) @@ -317,7 +317,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-computedarts") - dart_select = torch.cuda.DoubleTensor(total_overflow).uniform_() + dart_select = torch.cuda.FloatTensor(total_overflow).uniform_() # Compute darts for selection compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) @@ -408,7 +408,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # 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) + row_select_mtx_values = torch.cuda.FloatTensor(row_select_mtx_indices[0].size(0)).fill_(1.0) repeated_next_frontier = next_frontier_select.clone() scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() @@ -416,7 +416,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota col_select_mtx_rows = repeated_next_frontier.view(-1) col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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.FloatTensor(col_select_mtx_rows.size(0)).fill_(1.0) torch.cuda.nvtx.range_pop() timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 11c73a0..fa473ad 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -82,7 +82,7 @@ def scale_elements(adj_matrix, adj_part, node_count, row_vtx, col_vtx, normaliza return adj_part adj_part = adj_part.coalesce() - deg = torch.histc(adj_matrix[0].double(), bins=node_count) + deg = torch.histc(adj_matrix[0].float(), bins=node_count) deg = deg.pow(-0.5) row_len = adj_part.size(0) @@ -142,8 +142,8 @@ def row_normalize(mx): 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().double(), - mx._indices(), mx._values().double(), + mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values().float(), + mx._indices(), mx._values().float(), 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, size=(r_mat_inv.size(0), mx.size(1))) @@ -190,7 +190,7 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep 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(), + torch.ones(am_partitions[i].size(1)).float(), 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], \ @@ -358,7 +358,7 @@ def main(args): 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.DoubleTensor(batches_loc.size(1) * batches_loc.size(0)).fill_(1.0) + 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) torch.cuda.nvtx.range_pop() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 226dc8c..f7077dc 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -31,7 +31,7 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} -__device__ int binary_searchf(double *arr, double val, int imin, int imax) { +__device__ int binary_searchf(float *arr, float val, int imin, int imax) { int ans = 0; while (imax >= imin) { @@ -358,7 +358,7 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } -__global__ void ComputeDarts1D(double *dart_values, int n_darts, int mb_count) { +__global__ void ComputeDarts1D(float *dart_values, int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -366,7 +366,7 @@ __global__ void ComputeDarts1D(double *dart_values, int n_darts, int mb_count) { int dart_count = n_darts * mb_count; for (int i = id; i < dart_count; i += stride) { int row = i / n_darts; - dart_values[i] += (double)row; + dart_values[i] += (float)row; } } @@ -376,13 +376,13 @@ void compute_darts1d_gpu(const at::Tensor& dart_values, int n_darts, int mb_coun int BLOCK_COUNT = std::ceil((n_darts * mb_count) / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ComputeDarts1D<<>>(dart_values.data(), + ComputeDarts1D<<>>(dart_values.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, +__global__ void ComputeDartsSelect(float *dart_select, float *dart_hits_inv_sum, float *ps_dart_hits_inv_sum, long *ps_overflow, long mb_count, long total_overflow) { long id = blockIdx.x * blockDim.x + threadIdx.x; @@ -407,9 +407,9 @@ void compute_darts_select_gpu(const at::Tensor& dart_select, int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ComputeDartsSelect<<>>(dart_select.data(), - dart_hits_inv_sum.data(), - ps_dart_hits_inv_sum.data(), + ComputeDartsSelect<<>>(dart_select.data(), + dart_hits_inv_sum.data(), + ps_dart_hits_inv_sum.data(), ps_overflow.data(), mb_count, total_overflow); @@ -452,7 +452,7 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_values, int *h_map, +__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, int *h_values, int *h_map, int dart_count, int nnz) { int id = blockIdx.x * blockDim.x + threadIdx.x; @@ -480,8 +480,8 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, 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(), h_map.data(), dart_count, @@ -489,7 +489,7 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDartsSelect(double *dart_select, double *ps_dart_hits_inv, int *dart_hits_count, +__global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, int *dart_hits_count, int total_overflow, int nnz) { int id = blockIdx.x * blockDim.x + threadIdx.x; @@ -512,15 +512,15 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ThrowDartsSelect<<>>(dart_select.data(), - ps_dart_hits_inv.data(), + ThrowDartsSelect<<>>(dart_select.data(), + ps_dart_hits_inv.data(), dart_hits_count.data(), total_overflow, nnz); CHECK_ERROR("selection dart throwing error") } -__global__ void Normalize(double *output, double *input, long *index, int len) { +__global__ void Normalize(float *output, float *input, long *index, int len) { long id = blockIdx.x * blockDim.x + threadIdx.x; long stride = blockDim.x * gridDim.x; @@ -536,7 +536,7 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - Normalize<<>>(output.data(), input.data(), index.data(), len); + Normalize<<>>(output.data(), input.data(), index.data(), len); CHECK_ERROR("normalize error") } @@ -597,7 +597,7 @@ void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, i CHECK_ERROR("shift col select error") } -__global__ void ScatterAddD(double *src, long *indices, double *values, int num_vals) { +__global__ void ScatterAddD(float *src, long *indices, float *values, int num_vals) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -606,18 +606,18 @@ __global__ void ScatterAddD(double *src, long *indices, double *values, int num_ } } -void scatterd_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at::Tensor& values, int num_vals) { +void scatterf_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(), + ScatterAddD<<>>(src.data(), indices.data(), - values.data(), + values.data(), num_vals); - CHECK_ERROR("scatter add doubles error") + CHECK_ERROR("scatter add floats error") } __global__ void ScatterAddI(int *src, long *indices, int *values, int num_vals) { @@ -656,6 +656,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("normalize_gpu", &normalize_gpu, "Normalize values in an array based on a second array"); 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("scatterf_add_gpu", &scatterf_add_gpu, "Implementation of scatter_add_ for floats"); m.def("scatteri_add_gpu", &scatteri_add_gpu, "Implementation of scatter_add_ for ints"); } From a6b1d1fa69c3821a2c3216d8adbdb915c8e2e4f6 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 28 Jun 2022 17:03:32 -0700 Subject: [PATCH 078/178] Converted everything to doubles and timing_dict --- cagnet/samplers/ladies.py | 40 ++++++++++++------------ examples/gcn_15d.py | 8 ++--- sparse-extension/build.ou | 0 sparse-extension/sparse_coo_tensor.cu | 44 +++++++++++++-------------- 4 files changed, 47 insertions(+), 45 deletions(-) create mode 100644 sparse-extension/build.ou diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 2de72a0..7e0504f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -8,7 +8,7 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterf_add_gpu, scatteri_add_gpu + scatterd_add_gpu, scatteri_add_gpu def start_time(timer): timer.record() @@ -51,7 +51,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, matb_recv_values = matb._values().clone() else: matb_recv_indices = torch.cuda.LongTensor(2, matb_recv_nnz.item()) - matb_recv_values = torch.cuda.FloatTensor(matb_recv_nnz.item()) + matb_recv_values = torch.cuda.DoubleTensor(matb_recv_nnz.item()) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-inst-recv") start_time(start_timer) @@ -107,10 +107,10 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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.FloatTensor(matc_nnz).fill_(0.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.FloatTensor(matc_nnz - matc._nnz()).fill_(0.0))) + 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, "spgemm-padding") start_time(start_timer) @@ -183,19 +183,19 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota col_groups, timing_dict) torch.cuda.nvtx.range_pop() - print(f"probability-spgemm: {stop_time(start_timer, stop_timer)}") + timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p_den = torch.cuda.FloatTensor(mb_count).fill_(0) + p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) # p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) - scatterf_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) print(f"p.nnz: {p._nnz()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) - print(f"compute-p: {stop_time(start_timer, stop_timer)}") + timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") @@ -205,7 +205,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) torch.cuda.nvtx.range_pop() - print(f"pre-loop: {stop_time(start_timer, stop_timer)}") + timing_dict["pre-loop"].append(stop_time(start_timer, stop_timer)) iter_count = 0 selection_iter_count = 0 @@ -221,14 +221,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-prob-rowsum") - ps_p_values = torch.cumsum(p._values().float(), dim=0).roll(1) + ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) ps_p_values[0] = 0 torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") - dart_values = torch.cuda.FloatTensor(n_darts * mb_count).uniform_() + dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) @@ -289,7 +289,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-reciprocal") - dart_hits_inv = dart_hits_count.reciprocal().float() + 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)) @@ -303,9 +303,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-invsum") # dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() - dart_hits_inv_sum = torch.cuda.FloatTensor(mb_count).fill_(0) + dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) - scatterf_add_gpu(dart_hits_inv_sum, next_frontier._indices()[0, :], dart_hits_inv, \ + scatterd_add_gpu(dart_hits_inv_sum, next_frontier._indices()[0, :], dart_hits_inv, \ dart_hits_inv.size(0)) timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) @@ -317,7 +317,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-computedarts") - dart_select = torch.cuda.FloatTensor(total_overflow).uniform_() + dart_select = torch.cuda.DoubleTensor(total_overflow).uniform_() # Compute darts for selection compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) @@ -394,7 +394,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches._values().bool()).view(mb_count, batch_size) next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) torch.cuda.nvtx.range_pop() - print(f"construct-nextf: {stop_time(start_timer, stop_timer)}") + timing_dict["construct-nextf"].append(stop_time(start_timer, stop_timer)) torch.cuda.nvtx.range_push("nvtx-select-rowcols") # 1. Make the row/col select matrices @@ -408,7 +408,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO i > 0 (make sure size args to spspmm are right) print("ERROR i > 0") - row_select_mtx_values = torch.cuda.FloatTensor(row_select_mtx_indices[0].size(0)).fill_(1.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() scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() @@ -416,7 +416,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota col_select_mtx_rows = repeated_next_frontier.view(-1) col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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.FloatTensor(col_select_mtx_rows.size(0)).fill_(1.0) + col_select_mtx_values = torch.cuda.DoubleTensor(col_select_mtx_rows.size(0)).fill_(1.0) torch.cuda.nvtx.range_pop() timing_dict["select-mtxs"].append(stop_time(start_timer, stop_timer)) @@ -425,8 +425,10 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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)) + print(f"before row selection") sampled_indices, sampled_values = dist_spgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ row_groups, col_groups) + print(f"after row selection") sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, size=(nnz * mb_count, node_count_total)) @@ -470,7 +472,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) for k, v in timing_dict.items(): - print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)}") + print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)} len: {len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") return batches_select, next_frontier_select, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index fa473ad..c485732 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -142,8 +142,8 @@ def row_normalize(mx): 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().float(), - mx._indices(), mx._values().float(), + mx_indices, mx_values = torch_sparse.spspmm(r_mat_inv._indices(), r_mat_inv._values().double(), + mx._indices(), mx._values().double(), 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, size=(r_mat_inv.size(0), mx.size(1))) @@ -190,7 +190,7 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep 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)).float(), + 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], \ @@ -358,7 +358,7 @@ def main(args): 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_values = torch.cuda.DoubleTensor(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) torch.cuda.nvtx.range_pop() diff --git a/sparse-extension/build.ou b/sparse-extension/build.ou new file mode 100644 index 0000000..e69de29 diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index f7077dc..226dc8c 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -31,7 +31,7 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} -__device__ int binary_searchf(float *arr, float val, int imin, int imax) { +__device__ int binary_searchf(double *arr, double val, int imin, int imax) { int ans = 0; while (imax >= imin) { @@ -358,7 +358,7 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } -__global__ void ComputeDarts1D(float *dart_values, int n_darts, int mb_count) { +__global__ void ComputeDarts1D(double *dart_values, int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -366,7 +366,7 @@ __global__ void ComputeDarts1D(float *dart_values, int n_darts, int mb_count) { int dart_count = n_darts * mb_count; for (int i = id; i < dart_count; i += stride) { int row = i / n_darts; - dart_values[i] += (float)row; + dart_values[i] += (double)row; } } @@ -376,13 +376,13 @@ void compute_darts1d_gpu(const at::Tensor& dart_values, int n_darts, int mb_coun int BLOCK_COUNT = std::ceil((n_darts * mb_count) / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ComputeDarts1D<<>>(dart_values.data(), + ComputeDarts1D<<>>(dart_values.data(), n_darts, mb_count); CHECK_ERROR("dart1d computation error") } -__global__ void ComputeDartsSelect(float *dart_select, float *dart_hits_inv_sum, float *ps_dart_hits_inv_sum, +__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; @@ -407,9 +407,9 @@ void compute_darts_select_gpu(const at::Tensor& dart_select, int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ComputeDartsSelect<<>>(dart_select.data(), - dart_hits_inv_sum.data(), - ps_dart_hits_inv_sum.data(), + ComputeDartsSelect<<>>(dart_select.data(), + dart_hits_inv_sum.data(), + ps_dart_hits_inv_sum.data(), ps_overflow.data(), mb_count, total_overflow); @@ -452,7 +452,7 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, int *h_values, int *h_map, +__global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_values, int *h_map, int dart_count, int nnz) { int id = blockIdx.x * blockDim.x + threadIdx.x; @@ -480,8 +480,8 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, 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(), h_map.data(), dart_count, @@ -489,7 +489,7 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDartsSelect(float *dart_select, float *ps_dart_hits_inv, int *dart_hits_count, +__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; @@ -512,15 +512,15 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, int BLOCK_COUNT = std::ceil(total_overflow / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ThrowDartsSelect<<>>(dart_select.data(), - ps_dart_hits_inv.data(), + ThrowDartsSelect<<>>(dart_select.data(), + ps_dart_hits_inv.data(), dart_hits_count.data(), total_overflow, nnz); CHECK_ERROR("selection dart throwing error") } -__global__ void Normalize(float *output, float *input, long *index, int 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; @@ -536,7 +536,7 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - Normalize<<>>(output.data(), input.data(), index.data(), len); + Normalize<<>>(output.data(), input.data(), index.data(), len); CHECK_ERROR("normalize error") } @@ -597,7 +597,7 @@ void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, i CHECK_ERROR("shift col select error") } -__global__ void ScatterAddD(float *src, long *indices, float *values, int num_vals) { +__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; @@ -606,18 +606,18 @@ __global__ void ScatterAddD(float *src, long *indices, float *values, int num_va } } -void scatterf_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at::Tensor& values, int num_vals) { +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(), + ScatterAddD<<>>(src.data(), indices.data(), - values.data(), + values.data(), num_vals); - CHECK_ERROR("scatter add floats error") + CHECK_ERROR("scatter add doubles error") } __global__ void ScatterAddI(int *src, long *indices, int *values, int num_vals) { @@ -656,6 +656,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("normalize_gpu", &normalize_gpu, "Normalize values in an array based on a second array"); 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("scatterf_add_gpu", &scatterf_add_gpu, "Implementation of scatter_add_ for floats"); + 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"); } From 830ee5e2cfb8da440813fd08c98f8f5c2fbc1221 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 29 Jun 2022 00:27:20 -0700 Subject: [PATCH 079/178] Merged some spgemm timers into spgemm-misc --- cagnet/samplers/ladies.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 7e0504f..bcd3cf1 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -20,7 +20,15 @@ def stop_time(start_timer, stop_timer): def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if timing_dict is not None: - timing_dict[range_name].append(stop_time(start_timer, stop_timer)) + misc_range = range_name.startswith("spgemm-") and range_name != "spgemm-bcast-nnz" and \ + range_name != "spgemm-bcast-data" and \ + range_name != "spgemm-local-spgemm" and \ + range_name != "spgemm-reduce-nnz" and \ + range_name != "spgemm-reduce" + if misc_range: + timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) + else: + timing_dict[range_name].append(stop_time(start_timer, stop_timer)) def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ timing_dict=None): @@ -126,8 +134,11 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, start_time(start_timer) matc_recv = torch.stack(matc_recv) - matc = torch.sparse.sum(matc_recv, dim=0).coalesce() + matc = torch.sparse.sum(matc_recv, dim=0) stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-reduce") + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-reduce-coalesce") start_time(start_timer) nnz_mask = matc._values() != -1.0 @@ -425,10 +436,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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)) - print(f"before row selection") sampled_indices, sampled_values = dist_spgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ - row_groups, col_groups) - print(f"after row selection") + row_groups, col_groups, timing_dict) sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, size=(nnz * mb_count, node_count_total)) @@ -453,7 +462,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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))) sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ - row_groups, col_groups) + row_groups, col_groups, timing_dict) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -473,6 +482,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) for k, v in timing_dict.items(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)} len: {len(v)}") + if k.startswith("spgemm") and k != "spgemm-misc": + print(f"{k} times: {v}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") return batches_select, next_frontier_select, adj_matrices From 22ec2b615817c58c8bec62832da1d00d7bde40cf Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 14 Jul 2022 14:06:10 -0700 Subject: [PATCH 080/178] Fixed binary search bounds --- .gitignore | 1 + cagnet/samplers/ladies.py | 77 ++++++++++++++++----------- sparse-extension/build.ou | 0 sparse-extension/sparse_coo_tensor.cu | 39 +++++++++++--- 4 files changed, 78 insertions(+), 39 deletions(-) delete mode 100644 sparse-extension/build.ou diff --git a/.gitignore b/.gitignore index 867478c..932b594 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ sparse-extension/dist/ sparse-extension/sparse_coo_tensor_cpp.egg-info/ sparse_coo_tensor_cpp.egg-info/ .*.swp +.*.swo .nfs* *.txt !requirements.txt diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index bcd3cf1..3a7685a 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -20,18 +20,18 @@ def stop_time(start_timer, stop_timer): def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if timing_dict is not None: - misc_range = range_name.startswith("spgemm-") and range_name != "spgemm-bcast-nnz" and \ - range_name != "spgemm-bcast-data" and \ - range_name != "spgemm-local-spgemm" and \ - range_name != "spgemm-reduce-nnz" and \ - range_name != "spgemm-reduce" + misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ + not range_name.startswith("spgemm-bcast-data") and \ + not range_name.startswith("spgemm-local-spgemm") and \ + not range_name.startswith("spgemm-reduce-nnz") and \ + not range_name.startswith("spgemm-reduce") if misc_range: timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) else: timing_dict[range_name].append(stop_time(start_timer, stop_timer)) def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ - timing_dict=None): + name, timing_dict=None): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -44,14 +44,14 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stages = size // (replication ** 2) if rank_col == replication - 1: stages = (size // replication) - (replication - 1) * stages - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-matc-inst") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-inst-{name}") 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, "spgemm-bcast-nnz") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-bcast-nnz-{name}") start_time(start_timer) if q == rank: @@ -60,12 +60,12 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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, "spgemm-inst-recv") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-inst-recv-{name}") 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, "spgemm-bcast-data") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-bcast-data-{name}") start_time(start_timer) am_partid = rank_col * (size // replication ** 2) + i @@ -77,30 +77,30 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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, "spgemm-preproc-local") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}") start_time(start_timer) 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) - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-local-spgemm") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") start_time(start_timer) matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-chunk-inst") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") start_time(start_timer) matc_chunk = matc_chunk.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-chunk-coalesce") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") start_time(start_timer) matc += matc_chunk - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-chunk-add") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") start_time(start_timer) matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-matc-coalesce") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) # Implement sparse allreduce w/ all_gather and padding @@ -108,7 +108,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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, "spgemm-reduce-nnz") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-nnz-{name}") start_time(start_timer) matc_recv_indices = [] @@ -119,32 +119,32 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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, "spgemm-padding") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-padding-{name}") 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, "spgemm-allgather") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-allgather-{name}") 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())) - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-preproc-reduce") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}") 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, "spgemm-reduce") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}") start_time(start_timer) matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-reduce-coalesce") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") start_time(start_timer) nnz_mask = matc._values() != -1.0 matc_nnz_indices = matc._indices()[:, nnz_mask] matc_nnz_values = matc._values()[nnz_mask] - stop_time_add(start_timer, stop_timer, timing_dict, "spgemm-unpad") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-unpad-{name}") return matc_nnz_indices, matc_nnz_values @@ -190,8 +190,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, row_groups, \ - col_groups, timing_dict) + p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, \ + row_groups, col_groups, "prob", timing_dict) torch.cuda.nvtx.range_pop() timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -227,6 +227,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota while underfull_minibatches: iter_count += 1 + print(f"iter_count: {iter_count}", flush=True) start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -253,8 +254,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-filter-darts") # dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) - dart_hits_map = torch.cuda.IntTensor(n_darts * mb_count) - throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, dart_hits_map, \ + print(f"ps_p.values.size: {ps_p_values.size()} p.nnz: {p._nnz()}") + throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) @@ -437,7 +438,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota row_select_mtx = torch.sparse_coo_tensor(row_select_mtx_indices, row_select_mtx_values, size=(nnz * mb_count, node_count_total)) sampled_indices, sampled_values = dist_spgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ - row_groups, col_groups, timing_dict) + row_groups, col_groups, "rowsel", timing_dict) sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, size=(nnz * mb_count, node_count_total)) @@ -462,7 +463,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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))) sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ - row_groups, col_groups, timing_dict) + row_groups, col_groups, "colsel", timing_dict) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -481,9 +482,23 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) for k, v in timing_dict.items(): - print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)} len: {len(v)}") if k.startswith("spgemm") and k != "spgemm-misc": - print(f"{k} times: {v}") + # print(f"{k} times: {v}") + v_tens = torch.cuda.FloatTensor([sum(v)]) + v_tens_recv = [] + for i in range(size): + v_tens_recv.append(torch.cuda.FloatTensor(1)) + 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}") + for k, v in timing_dict.items(): + print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)} len: {len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") return batches_select, next_frontier_select, adj_matrices diff --git a/sparse-extension/build.ou b/sparse-extension/build.ou deleted file mode 100644 index e69de29..0000000 diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 226dc8c..c84a3ff 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -48,6 +48,24 @@ __device__ int binary_searchf(double *arr, double val, int imin, int imax) { 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; @@ -452,26 +470,32 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_values, int *h_map, +__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); - if (vtx < 0 || vtx >= nnz) { - printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); - } + int vtx = -1; + if (i == 1954) { + vtx = binary_searchfpf(ps_p_values, dart_values[i], 0, nnz - 1); + if (vtx < 0 || vtx >= nnz) { + printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); + } + } else { + vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz - 1); + if (vtx < 0 || vtx >= nnz) { + printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); + } + } atomicAdd(&h_values[vtx], 1); - h_map[i] = vtx; } } void throw_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& ps_p_values, const at::Tensor& h_values, - const at::Tensor& h_map, int dart_count, int nnz) { @@ -483,7 +507,6 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, ThrowDarts1D<<>>(dart_values.data(), ps_p_values.data(), h_values.data(), - h_map.data(), dart_count, nnz); CHECK_ERROR("dart throwing error") From a085d279ddec8c974cede6779b5c2b7c661c3fab Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 14 Jul 2022 14:15:31 -0700 Subject: [PATCH 081/178] Cleanup --- cagnet/samplers/ladies.py | 1 - sparse-extension/sparse_coo_tensor.cu | 16 ++++------------ 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 3a7685a..5714cf6 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -254,7 +254,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-filter-darts") # dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) - print(f"ps_p.values.size: {ps_p_values.size()} p.nnz: {p._nnz()}") throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index c84a3ff..03480b1 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -477,18 +477,10 @@ __global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_va int stride = blockDim.x * gridDim.x; for (int i = id; i < dart_count; i += stride) { - int vtx = -1; - if (i == 1954) { - vtx = binary_searchfpf(ps_p_values, dart_values[i], 0, nnz - 1); - if (vtx < 0 || vtx >= nnz) { - printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); - } - } else { - vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz - 1); - if (vtx < 0 || vtx >= nnz) { - printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); - } - } + int vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz - 1); + if (vtx < 0 || vtx >= nnz) { + printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); + } atomicAdd(&h_values[vtx], 1); } } From 0a8a9d2606b263e56cfb922b4c1dce797433b7d3 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 25 Jul 2022 10:55:29 -0700 Subject: [PATCH 082/178] Experimenting with different dart counts --- cagnet/samplers/ladies.py | 4 +- examples/gcn_15d.py | 14 + examples/tmp2.out | 312 ----------------- examples/tmp3.out | 473 -------------------------- examples/tmp_bsearch_fpoint_error.out | 259 -------------- 5 files changed, 15 insertions(+), 1047 deletions(-) delete mode 100644 examples/tmp2.out delete mode 100644 examples/tmp3.out delete mode 100644 examples/tmp_bsearch_fpoint_error.out diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 5714cf6..8e38dc9 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -227,7 +227,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota while underfull_minibatches: iter_count += 1 - print(f"iter_count: {iter_count}", flush=True) start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") @@ -252,7 +251,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-filter-darts") - # dart_hits_count = torch.cuda.LongTensor(p._nnz()).fill_(0) dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ n_darts * mb_count, p._nnz()) @@ -481,7 +479,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) for k, v in timing_dict.items(): - if k.startswith("spgemm") and k != "spgemm-misc": + if (k.startswith("spgemm") and k != "spgemm-misc") or k == "probability-spgemm" or k == "row-select-spgemm" or k == "col-select-spgemm": # print(f"{k} times: {v}") v_tens = torch.cuda.FloatTensor([sum(v)]) v_tens_recv = [] diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index c485732..ddc2761 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -364,10 +364,24 @@ def main(args): torch.cuda.nvtx.range_pop() node_count = inputs.size(0) + adj_matrix = adj_matrix.cuda() if args.n_darts == -1: edge_count = adj_matrix.size(1) avg_degree = int(edge_count / node_count) args.n_darts = avg_degree * args.batch_size + + degs = torch.cuda.IntTensor(node_count).fill_(0) + degs.scatter_add_(0, adj_matrix[0, :].long(), torch.cuda.IntTensor(edge_count).fill_(1)) + degs = degs[degs.nonzero()] + print(f"degs: {degs.squeeze()}") + max_degree = torch.max(degs).item() + med_degree = torch.median(degs).item() + print(f"degs.size: {degs.size()}") + print(f"degs.nnz.size: {degs.nonzero().squeeze().size()}") + + print(f"edge_count: {edge_count} node_count: {node_count} avg_degree: {avg_degree} max_degree: {max_degree} med_degree: {med_degree}") + + args.n_darts = int(max_degree / avg_degree) print(f"n_darts: {args.n_darts}") print(f"rank: {rank} g_loc: {g_loc}") diff --git a/examples/tmp2.out b/examples/tmp2.out deleted file mode 100644 index aea3f83..0000000 --- a/examples/tmp2.out +++ /dev/null @@ -1,312 +0,0 @@ -0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=1, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) -0: hostname: nid002964 -0: hostname: nid002964 rank: 0 size: 1 -0: start partitioning -0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) -0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) -0: end partitioning -0: coalescing -0: normalizing -0: done normalizing -0: n_darts: 125952 -0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], -0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), -0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), -0: device='cuda:0', size=(232965, 232965), nnz=114848857, -0: dtype=torch.float64, layout=torch.sparse_coo) -0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0, 0, 0, 0, -0: 0, 0, 0, 0], -0: [195991, 117560, 100852, 197262, 114239, 24469, 97172, -0: 12309, 212877, 77891, 135178, 163418, 60833, 125788, -0: 122201, 165750, 106560, 62024, 203033, 198683, 118941, -0: 2313, 90070, 217127, 153723, 220936, 106445, 76218, -0: 201129, 106587, 110028, 193334, 20554, 79072, 127722, -0: 137651, 214143, 177731, 30765, 91760, 101839, 39926, -0: 194008, 9498, 190477, 212081, 179417, 135134, 203107, -0: 230411, 8706, 185588, 52276, 184932, 61362, 95713, -0: 80064, 175943, 221178, 119857, 129676, 10259, 174825, -0: 208831, 61946, 99332, 214464, 184739, 49674, 164221, -0: 59486, 177651, 77796, 85412, 13951, 48978, 62602, -0: 225615, 182931, 200535, 207967, 103437, 70222, 223332, -0: 85366, 157946, 192683, 23053, 92330, 197448, 15910, -0: 19101, 171437, 179931, 52280, 197172, 60523, 154098, -0: 186386, 201681, 193268, 80456, 201243, 11949, 33508, -0: 19455, 115980, 46319, 39589, 170252, 174478, 97713, -0: 163749, 148905, 144319, 11441, 54393, 61999, 100190, -0: 83454, 114559, 59890, 39951, 88308, 88733, 142830, -0: 72414, 32648, 114384, 126907, 96293, 197589, 165314, -0: 55684, 120097, 134402, 59885, 114336, 113777, 36989, -0: 192085, 169890, 5813, 31575, 230457, 20423, 107725, -0: 77745, 13674, 13136, 12013, 213858, 97307, 180094, -0: 5329, 9922, 55727, 83586, 221433, 133085, 125773, -0: 209438, 151068, 79892, 137760, 232603, 57682, 110255, -0: 37831, 86141, 56670, 87840, 146715, 147534, 150393, -0: 58304, 206360, 56253, 196089, 48612, 28102, 53804, -0: 207564, 161368, 192228, 24247, 84217, 120440, 54570, -0: 159631, 93878, 1928, 169073, 122680, 135211, 6105, -0: 12968, 66029, 142093, 149634, 178465, 34053, 58560, -0: 161657, 130137, 17574, 137016, 80219, 176259, 101294, -0: 88046, 42065, 199930, 120417, 75237, 71001, 114135, -0: 110924, 20081, 207091, 64444, 105467, 163630, 195228, -0: 7255, 139108, 89413, 227103, 94165, 56486, 228263, -0: 223212, 214386, 32605, 75837, 73859, 67370, 125235, -0: 8321, 116827, 168590, 8279, 152684, 145955, 26546, -0: 137981, 13286, 77317, 25850, 83672, 84947, 155106, -0: 36568, 15526, 216693, 74749]]), -0: values=tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., -0: 1., 1., 1., 1.]), -0: device='cuda:0', size=(1, 232965), nnz=256, dtype=torch.float64, -0: layout=torch.sparse_coo) -0: probability-spgemm: 339.2951354980469 -0: before p._values(): tensor([ 7.8125e-03, 7.8125e-03, 7.8125e-03, ..., 3.6065e-319, -0: 3.6065e-319, 3.6066e-319], device='cuda:0', dtype=torch.float64) -0: compute-p: 1.1612800359725952 -0: p.nnz: 72999 -0: pre-loop: 0.19340799748897552 -0: imid: 36499 val: 0.877743 arr[imid]: 1.000000 -0: imid: 18249 val: 0.877743 arr[imid]: 0.452343 -0: imid: 27374 val: 0.877743 arr[imid]: 0.536156 -0: imid: 31936 val: 0.877743 arr[imid]: 0.795653 -0: imid: 34217 val: 0.877743 arr[imid]: 0.974655 -0: imid: 33076 val: 0.877743 arr[imid]: 0.917219 -0: imid: 32506 val: 0.877743 arr[imid]: 0.841953 -0: imid: 32791 val: 0.877743 arr[imid]: 0.890303 -0: imid: 32648 val: 0.877743 arr[imid]: 0.843019 -0: imid: 32719 val: 0.877743 arr[imid]: 0.890276 -0: imid: 32683 val: 0.877743 arr[imid]: 0.890252 -0: imid: 32665 val: 0.877743 arr[imid]: 0.890068 -0: imid: 32656 val: 0.877743 arr[imid]: 0.890068 -0: imid: 32652 val: 0.877743 arr[imid]: 0.843019 -0: imid: 32654 val: 0.877743 arr[imid]: 0.843019 -0: imid: 32655 val: 0.877743 arr[imid]: 0.843019 -0: iter_count: 1 dart_values: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', -0: dtype=torch.float64) -0: iter_count: 1 dart_values_scaled: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', -0: dtype=torch.float64) -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 16], -0: [ 58], -0: [ 87], -0: ..., -0: [36389], -0: [36423], -0: [36479]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([1036, 1]) -0: iter_count: 1 dart_hits_count.sum: 125952 -0: iter_count: 1 dart_hits_map: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', -0: dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[1], -0: [1], -0: [6], -0: ..., -0: [6], -0: [1], -0: [2]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.3693e-12, 2.7386e-12, ..., 1.0000e+00, 1.0000e+00, -0: 1.0000e+00], device='cuda:0', dtype=torch.float64) -0: construct-nextf: 0.2680000066757202 -0: total_time: 389.8416748046875 -0: imid: 18250 val: 0.877743 arr[imid]: 0.452343 -0: imid: 27375 val: 0.877743 arr[imid]: 0.536156 -0: imid: 31938 val: 0.877743 arr[imid]: 0.795653 -0: imid: 34219 val: 0.877743 arr[imid]: 0.974656 -0: imid: 33078 val: 0.877743 arr[imid]: 0.917219 -0: imid: 32508 val: 0.877743 arr[imid]: 0.841953 -0: imid: 32793 val: 0.877743 arr[imid]: 0.890303 -0: imid: 32650 val: 0.877743 arr[imid]: 0.843019 -0: imid: 32721 val: 0.877743 arr[imid]: 0.890276 -0: imid: 32685 val: 0.877743 arr[imid]: 0.890252 -0: imid: 32667 val: 0.877743 arr[imid]: 0.890252 -0: imid: 32658 val: 0.877743 arr[imid]: 0.890068 -0: imid: 32654 val: 0.877743 arr[imid]: 0.843019 -0: imid: 32656 val: 0.877743 arr[imid]: 0.890068 -0: imid: 32655 val: 0.877743 arr[imid]: 0.843019 -0: spgemm-matc-inst total_time: 0.06140799820423126 avg_time 0.06140799820423126 -0: spgemm-bcast-nnz total_time: 123.69152069091797 avg_time 123.69152069091797 -0: spgemm-inst-recv total_time: 4.317759990692139 avg_time 4.317759990692139 -0: spgemm-bcast-data total_time: 0.09142400324344635 avg_time 0.09142400324344635 -0: spgemm-preproc-local total_time: 0.436352014541626 avg_time 0.436352014541626 -0: spgemm-local-spgemm total_time: 170.251708984375 avg_time 170.251708984375 -0: spgemm-chunk-inst total_time: 0.17958399653434753 avg_time 0.17958399653434753 -0: spgemm-chunk-coalesce total_time: 0.49852800369262695 avg_time 0.49852800369262695 -0: spgemm-chunk-add total_time: 0.11667200177907944 avg_time 0.11667200177907944 -0: spgemm-matc-coalesce total_time: 0.011359999887645245 avg_time 0.011359999887645245 -0: spgemm-reduce-nnz total_time: 37.1005744934082 avg_time 37.1005744934082 -0: spgemm-padding total_time: 0.14291200041770935 avg_time 0.14291200041770935 -0: spgemm-allgather total_time: 0.2025279998779297 avg_time 0.2025279998779297 -0: spgemm-preproc-reduce total_time: 0.15545600652694702 avg_time 0.15545600652694702 -0: spgemm-reduce total_time: 1.304703950881958 avg_time 1.304703950881958 -0: spgemm-unpad total_time: 0.26422399282455444 avg_time 0.26422399282455444 -0: prob-rowsum total_time: 0.1865919977426529 avg_time 0.1865919977426529 -0: gen-darts total_time: 0.0759039968252182 avg_time 0.0759039968252182 -0: dart-throw total_time: 1.1985280513763428 avg_time 1.1985280513763428 -0: filter_darts total_time: 2.5079360008239746 avg_time 2.5079360008239746 -0: add-to-frontier total_time: 0.27209600806236267 avg_time 0.27209600806236267 -0: count_samples total_time: 0.13497599959373474 avg_time 0.13497599959373474 -0: select-preproc total_time: 0.1653759926557541 avg_time 0.1653759926557541 -0: select-psoverflow total_time: 0.4081280007958412 avg_time 0.08162560015916824 -0: select-reciprocal total_time: 1.1649919897317886 avg_time 0.23299839794635774 -0: select-instmtx total_time: 0.729408010840416 avg_time 0.14588160216808319 -0: select-invsum total_time: 1.12848000228405 avg_time 0.22569600045681 -0: select-psinv total_time: 0.48076799511909485 avg_time 0.09615359902381897 -0: select-computedarts total_time: 0.30032000318169594 avg_time 0.06006400063633919 -0: select-throwdarts total_time: 0.3787840008735657 avg_time 0.07575680017471313 -0: select-add-to-frontier total_time: 0.9426880180835724 avg_time 0.18853760361671448 -0: select-samplecount total_time: 0.667136013507843 avg_time 0.1334272027015686 -0: select-overflow total_time: 0.4785599932074547 avg_time 0.09571199864149094 -0: select-iter total_time: 7.840160012245178 avg_time 1.5680320024490357 -0: dart-selection total_time: 8.04736042022705 avg_time 8.04736042022705 -0: set-probs total_time: 0.3381119966506958 avg_time 0.3381119966506958 -0: compute-bool total_time: 0.06675200164318085 avg_time 0.06675200164318085 -0: sampling-iters total_time: 13.418304443359375 avg_time 13.418304443359375 -0: select-mtxs total_time: 0.2803199887275696 avg_time 0.2803199887275696 -0: row-select-spgemm total_time: 20.787424087524414 avg_time 20.787424087524414 -0: row-select-expand total_time: 0.10172799974679947 avg_time 0.10172799974679947 -0: col-select-spgemm total_time: 13.91817569732666 avg_time 13.91817569732666 -0: set-sample total_time: 0.12230399996042252 avg_time 0.12230399996042252 -0: iter_count: 1 -0: selection_iter_count: 5 -0: -0: probability-spgemm: 52.053855895996094 -0: before p._values(): tensor([ 7.8125e-03, 7.8125e-03, 7.8125e-03, ..., 2.0000e+00, -0: 2.0000e+00, 5.2635e-315], device='cuda:0', dtype=torch.float64) -0: compute-p: 0.8244479894638062 -0: p.nnz: 36500 -0: pre-loop: 0.16601599752902985 -0: iter_count: 1 dart_values: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', -0: dtype=torch.float64) -0: iter_count: 1 dart_values_scaled: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', -0: dtype=torch.float64) -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 16], -0: [ 58], -0: [ 87], -0: ..., -0: [36389], -0: [36423], -0: [36479]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([1036, 1]) -0: iter_count: 1 dart_hits_count.sum: 125952 -0: iter_count: 1 dart_hits_map: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', -0: dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[1], -0: [1], -0: [6], -0: ..., -0: [6], -0: [1], -0: [2]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.3693e-12, 2.7386e-12, ..., 1.0000e+00, 1.0000e+00, -0: 1.0000e+00], device='cuda:0', dtype=torch.float64) -0: construct-nextf: 0.25760000944137573 -0: total_time: 100.39507293701172 -0: spgemm-matc-inst total_time: 0.039872001856565475 avg_time 0.039872001856565475 -0: spgemm-bcast-nnz total_time: 0.08102399855852127 avg_time 0.08102399855852127 -0: spgemm-inst-recv total_time: 4.348288059234619 avg_time 4.348288059234619 -0: spgemm-bcast-data total_time: 0.061664000153541565 avg_time 0.061664000153541565 -0: spgemm-preproc-local total_time: 0.28252801299095154 avg_time 0.28252801299095154 -0: spgemm-local-spgemm total_time: 44.76665496826172 avg_time 44.76665496826172 -0: spgemm-chunk-inst total_time: 0.14655999839305878 avg_time 0.14655999839305878 -0: spgemm-chunk-coalesce total_time: 0.4310719966888428 avg_time 0.4310719966888428 -0: spgemm-chunk-add total_time: 0.06966400146484375 avg_time 0.06966400146484375 -0: spgemm-matc-coalesce total_time: 0.012160000391304493 avg_time 0.012160000391304493 -0: spgemm-reduce-nnz total_time: 0.09187199920415878 avg_time 0.09187199920415878 -0: spgemm-padding total_time: 0.11664000153541565 avg_time 0.11664000153541565 -0: spgemm-allgather total_time: 0.1693439930677414 avg_time 0.1693439930677414 -0: spgemm-preproc-reduce total_time: 0.15148800611495972 avg_time 0.15148800611495972 -0: spgemm-reduce total_time: 0.6639999747276306 avg_time 0.6639999747276306 -0: spgemm-unpad total_time: 0.22323200106620789 avg_time 0.22323200106620789 -0: prob-rowsum total_time: 0.10636799782514572 avg_time 0.10636799782514572 -0: gen-darts total_time: 0.03984000161290169 avg_time 0.03984000161290169 -0: dart-throw total_time: 1.133247971534729 avg_time 1.133247971534729 -0: filter_darts total_time: 2.332927942276001 avg_time 2.332927942276001 -0: add-to-frontier total_time: 0.1860799938440323 avg_time 0.1860799938440323 -0: count_samples total_time: 0.10454399883747101 avg_time 0.10454399883747101 -0: select-preproc total_time: 0.09487999975681305 avg_time 0.09487999975681305 -0: select-psoverflow total_time: 0.39632000029087067 avg_time 0.07926400005817413 -0: select-reciprocal total_time: 0.8863359987735748 avg_time 0.17726719975471497 -0: select-instmtx total_time: 0.6157120019197464 avg_time 0.12314240038394927 -0: select-invsum total_time: 0.7136320024728775 avg_time 0.1427264004945755 -0: select-psinv total_time: 0.4678719937801361 avg_time 0.09357439875602722 -0: select-computedarts total_time: 0.30399999767541885 avg_time 0.06079999953508377 -0: select-throwdarts total_time: 0.35942399501800537 avg_time 0.07188479900360108 -0: select-add-to-frontier total_time: 0.8053120076656342 avg_time 0.16106240153312684 -0: select-samplecount total_time: 0.48419199883937836 avg_time 0.09683839976787567 -0: select-overflow total_time: 0.46035199612379074 avg_time 0.09207039922475815 -0: select-iter total_time: 6.5983359813690186 avg_time 1.3196671962738038 -0: dart-selection total_time: 6.795360088348389 avg_time 6.795360088348389 -0: set-probs total_time: 0.2848320007324219 avg_time 0.2848320007324219 -0: compute-bool total_time: 0.05846399813890457 avg_time 0.05846399813890457 -0: sampling-iters total_time: 11.527423858642578 avg_time 11.527423858642578 -0: select-mtxs total_time: 0.25360000133514404 avg_time 0.25360000133514404 -0: row-select-spgemm total_time: 20.890464782714844 avg_time 20.890464782714844 -0: row-select-expand total_time: 0.10073599964380264 avg_time 0.10073599964380264 -0: col-select-spgemm total_time: 13.9202880859375 avg_time 13.9202880859375 -0: set-sample total_time: 0.12569600343704224 avg_time 0.12569600343704224 -0: iter_count: 1 -0: selection_iter_count: 5 -0: -0: return_dart_count1.nnzidxs: tensor([ 16, 58, 87, ..., 36389, 36423, 36479], device='cuda:0') -0: return_dart_count1: tensor([1, 1, 6, ..., 6, 1, 2], device='cuda:0', dtype=torch.int32) -0: return_dart_map1: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', -0: dtype=torch.int32) -0: return_dart_val1: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', -0: dtype=torch.float64) -0: return_dart_count2.nnzidxs: tensor([ 16, 58, 87, ..., 36389, 36423, 36479], device='cuda:0') -0: return_dart_count2: tensor([1, 1, 6, ..., 6, 1, 2], device='cuda:0', dtype=torch.int32) -0: return_dart_map2: tensor([32656, 33660, 30747, ..., 21060, 32927, 30747], device='cuda:0', -0: dtype=torch.int32) -0: return_dart_val2: tensor([0.8805, 0.9397, 0.7129, ..., 0.4654, 0.8969, 0.7037], device='cuda:0', -0: dtype=torch.float64) -0: -0: unmatched_darts: tensor([], device='cuda:0', dtype=torch.int64) diff --git a/examples/tmp3.out b/examples/tmp3.out deleted file mode 100644 index d1c2b83..0000000 --- a/examples/tmp3.out +++ /dev/null @@ -1,473 +0,0 @@ -0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=512, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) -0: hostname: nid002781 -0: hostname: nid002781 rank: 0 size: 1 -0: start partitioning -0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) -0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) -0: end partitioning -0: coalescing -0: normalizing -0: done normalizing -0: imid: 18166615 val: 8172334.500000 arr[imid]: 2564784128.000000 -0: imid: 9083307 val: 8172334.500000 arr[imid]: 1288451712.000000 -0: imid: 4541653 val: 8172334.500000 arr[imid]: 640568896.000000 -0: imid: 2270826 val: 8172334.500000 arr[imid]: 317143936.000000 -0: imid: 1135412 val: 8172334.500000 arr[imid]: 159784208.000000 -0: imid: 567705 val: 8172334.500000 arr[imid]: 80147880.000000 -0: imid: 283852 val: 8172334.500000 arr[imid]: 41578028.000000 -0: imid: 141925 val: 8172334.500000 arr[imid]: 20473418.000000 -0: imid: 70962 val: 8172334.500000 arr[imid]: 9832408.000000 -0: imid: 35480 val: 8172334.500000 arr[imid]: 5203542.000000 -0: imid: 53221 val: 8172334.500000 arr[imid]: 7887794.000000 -0: imid: 62091 val: 8172334.500000 arr[imid]: 8652051.000000 -0: imid: 57656 val: 8172334.500000 arr[imid]: 8335342.500000 -0: imid: 55438 val: 8172334.500000 arr[imid]: 8064393.500000 -0: imid: 56547 val: 8172334.500000 arr[imid]: 8117769.000000 -0: imid: 57101 val: 8172334.500000 arr[imid]: 8215006.000000 -0: imid: 56824 val: 8172334.500000 arr[imid]: 8170865.500000 -0: imid: 56962 val: 8172334.500000 arr[imid]: 8189993.000000 -0: imid: 56893 val: 8172334.500000 arr[imid]: 8172020.500000 -0: imid: 56927 val: 8172334.500000 arr[imid]: 8188378.500000 -0: imid: 56910 val: 8172334.500000 arr[imid]: 8172235.500000 -0: imid: 56918 val: 8172334.500000 arr[imid]: 8188236.000000 -0: imid: 56914 val: 8172334.500000 arr[imid]: 8172336.000000 -0: imid: 56912 val: 8172334.500000 arr[imid]: 8172327.000000 -0: imid: 56913 val: 8172334.500000 arr[imid]: 8172329.000000 -0: imid: 18166615 val: 3401610.500000 arr[imid]: 0.000000 -0: imid: 27249923 val: 3401610.500000 arr[imid]: 0.000000 -0: imid: 31791577 val: 3401610.500000 arr[imid]: 9985493.000000 -0: imid: 29520750 val: 3401610.500000 arr[imid]: 0.000000 -0: imid: 30656163 val: 3401610.500000 arr[imid]: 0.000000 -0: imid: 31223870 val: 3401610.500000 arr[imid]: 9985493.000000 -0: imid: 30940016 val: 3401610.500000 arr[imid]: 0.000000 -0: imid: 31081943 val: 3401610.500000 arr[imid]: 0.000000 -0: imid: 31152906 val: 3401610.500000 arr[imid]: 6810933.000000 -0: imid: 31117424 val: 3401610.500000 arr[imid]: 2159459.750000 -0: imid: 31135165 val: 3401610.500000 arr[imid]: 4937014.500000 -0: imid: 31126294 val: 3401610.500000 arr[imid]: 3478861.000000 -0: imid: 31121859 val: 3401610.500000 arr[imid]: 2739246.250000 -0: imid: 31124076 val: 3401610.500000 arr[imid]: 3059827.500000 -0: imid: 31125185 val: 3401610.500000 arr[imid]: 3233514.500000 -0: imid: 31125739 val: 3401610.500000 arr[imid]: 3426718.250000 -0: imid: 31125462 val: 3401610.500000 arr[imid]: 3257614.000000 -0: imid: 31125600 val: 3401610.500000 arr[imid]: 3265417.500000 -0: imid: 31125669 val: 3401610.500000 arr[imid]: 3318082.250000 -0: imid: 31125704 val: 3401610.500000 arr[imid]: 3421715.250000 -0: imid: 31125686 val: 3401610.500000 arr[imid]: 3421321.000000 -0: imid: 31125677 val: 3401610.500000 arr[imid]: 3420803.250000 -0: imid: 31125673 val: 3401610.500000 arr[imid]: 3420798.750000 -0: imid: 31125671 val: 3401610.500000 arr[imid]: 3318083.750000 -0: imid: 31125672 val: 3401610.500000 arr[imid]: 3318098.000000 -0: n_darts: 125952 -0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], -0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), -0: values=tensor([2.0568e-07, 2.0568e-07, 2.0568e-07, ..., -0: 4.1311e-06, 4.1311e-06, 4.1311e-06]), -0: device='cuda:0', size=(232965, 232965), nnz=114848857, -0: layout=torch.sparse_coo) -0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 511, 511, 511], -0: [195991, 117560, 100852, ..., 210813, 53334, 131122]]), -0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), -0: device='cuda:0', size=(512, 232965), nnz=131072, layout=torch.sparse_coo) -0: probability-spgemm: 343.3492736816406 -0: before p._values(): tensor([1.2276e+00, 1.8491e-01, 1.8721e+01, ..., 3.6240e-01, 7.9862e+00, -0: 2.4029e+02], device='cuda:0') -0: before p._values()[56912]: 13.246072769165039 p_den[row[56912]]: 6.2924485206604 -0: compute-p: 61.61203384399414 -0: p.nnz: 36333230 -0: pre-loop: 1.4520319700241089 -0: p._values()[56912]: 2.105074405670166 ps_p_values[56912]: 8172327.0 -0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') -0: iter_count: 1 dart_values_scaled: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, -0: 5.1157e+09], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], -0: [ 8], -0: [ 11], -0: ..., -0: [36332465], -0: [36332476], -0: [36332500]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3563590, 1]) -0: iter_count: 1 dart_hits_count.sum: 64487424 -0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[ 1], -0: [ 1], -0: [ 1], -0: ..., -0: [ 9], -0: [ 5], -0: [351]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9509e-01, 2.2447e-01, ..., 5.1202e+09, 5.1202e+09, -0: 5.1202e+09], device='cuda:0') -0: sampled_count: tensor([11024, 9666, 12520, 10538, 12932, 11469, 11153, 10402, 12684, 9842, -0: 12438, 13403, 12459, 10115, 11264, 12120, 13021, 13453, 12467, 11809, -0: 11179, 12947, 11294, 10516, 10519, 11119, 10005, 10003, 9595, 11595, -0: 13284, 11668, 8780, 9898, 12276, 11645, 10163, 10876, 10062, 9497, -0: 10663, 11791, 11570, 12157, 13784, 9359, 11432, 9263, 10236, 10962, -0: 9661, 10819, 10947, 10826, 10793, 10124, 8664, 9424, 10931, 9277, -0: 8969, 10122, 9018, 9295, 10272, 9501, 9365, 10336, 10197, 10276, -0: 8096, 10126, 9304, 9791, 9513, 8806, 8005, 9382, 7457, 8782, -0: 9591, 11029, 9743, 11499, 8958, 8565, 10265, 11149, 10869, 12021, -0: 8819, 9657, 7648, 9854, 10864, 9005, 9849, 9635, 10688, 9500, -0: 10095, 8867, 10907, 11618, 7965, 7262, 9538, 8247, 8774, 7919, -0: 7761, 7956, 7381, 8235, 8723, 9805, 8850, 7419, 8473, 7790, -0: 8735, 9197, 7685, 7017, 8090, 8958, 8467, 9538, 6961, 6832, -0: 9487, 8197, 8245, 7300, 8969, 7602, 7694, 8738, 7283, 7328, -0: 6346, 8730, 8239, 7295, 7483, 8942, 8718, 7367, 7956, 7476, -0: 7024, 7648, 7110, 9080, 8088, 9321, 8696, 8877, 7566, 9046, -0: 8364, 5883, 7471, 8178, 8085, 7091, 8098, 7641, 6458, 7959, -0: 7628, 9135, 9614, 6488, 8197, 7765, 10026, 6966, 7791, 7227, -0: 8740, 7021, 8104, 7890, 7465, 6500, 8742, 6896, 8484, 8183, -0: 8265, 7066, 7550, 6762, 7301, 7924, 7399, 6960, 6635, 8193, -0: 6658, 8946, 9388, 7317, 6856, 5976, 8725, 6323, 8338, 10963, -0: 8423, 6646, 8273, 8477, 7008, 6093, 5807, 5955, 5382, 5805, -0: 7292, 5586, 4903, 5398, 6756, 5422, 5375, 6447, 7429, 8104, -0: 7143, 7807, 6924, 6717, 6088, 5282, 7397, 5732, 5386, 7286, -0: 6578, 6753, 5975, 5473, 5728, 5789, 7533, 6489, 4791, 6969, -0: 6134, 6345, 5086, 5451, 6676, 5303, 5671, 5492, 5682, 5190, -0: 5790, 6892, 4640, 5092, 6053, 4556, 4478, 5817, 6266, 5773, -0: 5562, 5093, 7647, 5904, 5386, 7278, 5140, 4420, 5586, 5943, -0: 5722, 5877, 5422, 5381, 5198, 4538, 5971, 5915, 4892, 5487, -0: 5812, 5632, 5510, 7292, 5447, 4926, 5661, 5477, 6372, 5075, -0: 4935, 4570, 4551, 7161, 7198, 4613, 6546, 5083, 5790, 6444, -0: 6488, 5273, 5317, 6100, 6289, 5295, 5661, 5739, 4919, 6473, -0: 7321, 4467, 5162, 5626, 5163, 5242, 5738, 5566, 5668, 5366, -0: 5975, 6558, 4660, 5576, 4756, 5136, 5945, 5684, 5255, 5117, -0: 6296, 6262, 5105, 7194, 6323, 5272, 5927, 5982, 6396, 6069, -0: 6153, 7047, 6407, 5465, 6407, 5350, 5802, 5449, 5421, 6051, -0: 6494, 6544, 5567, 6552, 5296, 6268, 6327, 5792, 5634, 6637, -0: 5647, 5874, 5290, 4723, 4762, 5722, 5901, 5591, 5204, 6419, -0: 6053, 5675, 4983, 7242, 6072, 5344, 5710, 5781, 6664, 6667, -0: 5141, 5986, 4950, 4950, 4962, 6081, 6078, 4958, 5320, 5592, -0: 6145, 4617, 5423, 6351, 4915, 5882, 5784, 6182, 5397, 4903, -0: 5299, 6336, 5898, 6069, 4632, 5133, 6352, 5503, 5518, 5475, -0: 5775, 5268, 6240, 5027, 5523, 5914, 6114, 6852, 6178, 4491, -0: 4671, 4882, 4856, 4035, 4366, 4375, 4475, 4337, 4132, 4379, -0: 3563, 3770, 4201, 5074, 4138, 4800, 4169, 4887, 5110, 4784, -0: 3523, 4548, 4785, 3886, 4104, 4618, 4601, 4350, 4021, 3934, -0: 3695, 3647, 4122, 3875, 4670, 3810, 4459, 3030, 4273, 3011, -0: 4436, 4376, 4015, 4285, 4554, 4830, 4837, 5278, 4474, 4237, -0: 4983, 4089, 4022, 4586, 4258, 3467, 3603, 3786, 4418, 3655, -0: 3587, 3477, 3542, 5781, 3738, 4358, 3541, 3870, 3507, 3611, -0: 3784, 3311, 4182, 4893, 4322, 3625, 3868, 4655, 4744, 3510, -0: 4002, 3945], device='cuda:0', dtype=torch.int32) -0: p._values()[56912]: 0.0 ps_p_values[56912]: 0.0 -0: iter_count: 2 dart_values: tensor([0.4230, 0.5678, 0.2647, ..., 0.0191, 0.5965, 0.7411], device='cuda:0') -0: iter_count: 2 dart_values_scaled: tensor([4.2305e+06, 5.6784e+06, 2.6466e+06, ..., 5.1102e+09, 5.1160e+09, -0: 5.1174e+09], device='cuda:0') -0: iter_count: 2 dart_hits_count.nnzidxs: tensor([[ 0], -0: [31104767], -0: [31104773], -0: ..., -0: [31172045], -0: [31172053], -0: [31172057]], device='cuda:0') -0: iter_count: 2 dart_hits_count.nnzidxs.size: torch.Size([11976, 1]) -0: iter_count: 2 dart_hits_count.sum: 64487424 -0: iter_count: 2 dart_hits_map: tensor([31129386, 31143585, 31121122, ..., 0, 0, 0], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 2 dart_hits_count: tensor([[64361680], -0: [ 1], -0: [ 1], -0: ..., -0: [ 1], -0: [ 1], -0: [ 1]], device='cuda:0', dtype=torch.int32) -0: iter_count: 2 ps_p_values: tensor([ 0., 0., 0., ..., 9985493., 9985493., 9985493.], -0: device='cuda:0') -0: imid: 18166615 val: 8172334.500000 arr[imid]: 2564784896.000000 -0: imid: 9083307 val: 8172334.500000 arr[imid]: 1288451840.000000 -0: imid: 4541653 val: 8172334.500000 arr[imid]: 640569024.000000 -0: imid: 2270826 val: 8172334.500000 arr[imid]: 317144224.000000 -0: imid: 1135412 val: 8172334.500000 arr[imid]: 159784624.000000 -0: imid: 567705 val: 8172334.500000 arr[imid]: 80148176.000000 -0: imid: 283852 val: 8172334.500000 arr[imid]: 41578188.000000 -0: imid: 141925 val: 8172334.500000 arr[imid]: 20473426.000000 -0: imid: 70962 val: 8172334.500000 arr[imid]: 9832383.000000 -0: imid: 35480 val: 8172334.500000 arr[imid]: 5203529.000000 -0: imid: 53221 val: 8172334.500000 arr[imid]: 7887774.500000 -0: imid: 62091 val: 8172334.500000 arr[imid]: 8652029.000000 -0: imid: 57656 val: 8172334.500000 arr[imid]: 8335322.500000 -0: imid: 55438 val: 8172334.500000 arr[imid]: 8064373.500000 -0: imid: 56547 val: 8172334.500000 arr[imid]: 8117749.500000 -0: imid: 57101 val: 8172334.500000 arr[imid]: 8214986.500000 -0: imid: 56824 val: 8172334.500000 arr[imid]: 8170845.500000 -0: imid: 56962 val: 8172334.500000 arr[imid]: 8189973.500000 -0: imid: 56893 val: 8172334.500000 arr[imid]: 8172001.000000 -0: imid: 56927 val: 8172334.500000 arr[imid]: 8188358.500000 -0: imid: 56910 val: 8172334.500000 arr[imid]: 8172215.500000 -0: imid: 56918 val: 8172334.500000 arr[imid]: 8188216.000000 -0: imid: 56914 val: 8172334.500000 arr[imid]: 8172316.000000 -0: imid: 56916 val: 8172334.500000 arr[imid]: 8172323.500000 -0: imid: 56917 val: 8172334.500000 arr[imid]: 8188216.000000 -0: sampled_count: tensor([ 257, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 11990, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256, 256, 256, 256, 256, 256, 256, 256, 256, -0: 256, 256], device='cuda:0', dtype=torch.int32) -0: construct-nextf: 0.7972480058670044 -0: total_time: 3895.928466796875 -0: spgemm-matc-inst total_time: 0.05878400057554245 avg_time 0.05878400057554245 -0: spgemm-bcast-nnz total_time: 122.26163482666016 avg_time 122.26163482666016 -0: spgemm-inst-recv total_time: 3.5645439624786377 avg_time 3.5645439624786377 -0: spgemm-bcast-data total_time: 0.09523200243711472 avg_time 0.09523200243711472 -0: spgemm-preproc-local total_time: 0.3563520014286041 avg_time 0.3563520014286041 -0: spgemm-local-spgemm total_time: 96.65433502197266 avg_time 96.65433502197266 -0: spgemm-chunk-inst total_time: 1.1634880304336548 avg_time 1.1634880304336548 -0: spgemm-chunk-coalesce total_time: 33.33251190185547 avg_time 33.33251190185547 -0: spgemm-chunk-add total_time: 1.1100159883499146 avg_time 1.1100159883499146 -0: spgemm-matc-coalesce total_time: 0.01228800043463707 avg_time 0.01228800043463707 -0: spgemm-reduce-nnz total_time: 36.91321563720703 avg_time 36.91321563720703 -0: spgemm-padding total_time: 1.8552320003509521 avg_time 1.8552320003509521 -0: spgemm-allgather total_time: 2.268160104751587 avg_time 2.268160104751587 -0: spgemm-preproc-reduce total_time: 1.1533440351486206 avg_time 1.1533440351486206 -0: spgemm-reduce total_time: 39.05833435058594 avg_time 39.05833435058594 -0: spgemm-unpad total_time: 3.018752098083496 avg_time 3.018752098083496 -0: prob-rowsum total_time: 1.2051519751548767 avg_time 0.6025759875774384 -0: gen-darts total_time: 0.44383999705314636 avg_time 0.22191999852657318 -0: dart-throw total_time: 3.201184034347534 avg_time 1.600592017173767 -0: filter_darts total_time: 73.3745288848877 avg_time 36.68726444244385 -0: add-to-frontier total_time: 3.6911680698394775 avg_time 1.8455840349197388 -0: count_samples total_time: 61.86393737792969 avg_time 30.931968688964844 -0: select-preproc total_time: 0.2099199965596199 avg_time 0.10495999827980995 -0: select-psoverflow total_time: 1.9985919930040836 avg_time 0.062455999781377614 -0: select-reciprocal total_time: 41.37455999851227 avg_time 1.2929549999535084 -0: select-instmtx total_time: 36.437375664711 avg_time 1.1386679895222187 -0: select-invsum total_time: 1894.122688293457 avg_time 59.19133400917053 -0: select-psinv total_time: 2.2108160331845284 avg_time 0.06908800103701651 -0: select-computedarts total_time: 1.884160004556179 avg_time 0.058880000142380595 -0: select-throwdarts total_time: 10.209279954433441 avg_time 0.31903999857604504 -0: select-add-to-frontier total_time: 53.809184312820435 avg_time 1.6815370097756386 -0: select-samplecount total_time: 991.8761978149414 avg_time 30.99613118171692 -0: select-overflow total_time: 2.233344040811062 avg_time 0.06979200127534568 -0: select-iter total_time: 3043.3382110595703 avg_time 95.10431909561157 -0: dart-selection total_time: 3054.8369750976562 avg_time 1527.4184875488281 -0: set-probs total_time: 3.370400071144104 avg_time 1.685200035572052 -0: compute-bool total_time: 0.10444800183176994 avg_time 0.05222400091588497 -0: sampling-iters total_time: 3206.3609619140625 avg_time 1603.1804809570312 -0: select-mtxs total_time: 0.9584640264511108 avg_time 0.9584640264511108 -0: row-select-spgemm total_time: 198.93040466308594 avg_time 198.93040466308594 -0: row-select-expand total_time: 2.83516788482666 avg_time 2.83516788482666 -0: col-select-spgemm total_time: 79.16544342041016 avg_time 79.16544342041016 -0: set-sample total_time: 0.12639999389648438 avg_time 0.12639999389648438 -0: iter_count: 2 -0: selection_iter_count: 32 -0: -0: probability-spgemm: 184.3711700439453 -0: before p._values(): tensor([1.2276e+00, 1.8491e-01, 1.8721e+01, ..., 3.6240e-01, 7.9862e+00, -0: 2.4029e+02], device='cuda:0') -0: before p._values()[56912]: 13.246072769165039 p_den[row[56912]]: 6.292463779449463 -0: compute-p: 61.37958526611328 -0: p.nnz: 36333230 -0: pre-loop: 1.4438400268554688 -0: p._values()[56912]: 2.105069398880005 ps_p_values[56912]: 8172307.0 -0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') -0: iter_count: 1 dart_values_scaled: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, -0: 5.1157e+09], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], -0: [ 8], -0: [ 11], -0: ..., -0: [36332465], -0: [36332476], -0: [36332500]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3563119, 1]) -0: iter_count: 1 dart_hits_count.sum: 64487424 -0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[ 1], -0: [ 1], -0: [ 1], -0: ..., -0: [ 9], -0: [ 6], -0: [304]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9508e-01, 2.2447e-01, ..., 5.1202e+09, 5.1202e+09, -0: 5.1202e+09], device='cuda:0') -0: sampled_count: tensor([11038, 9667, 12520, 10568, 12768, 11472, 11310, 10395, 12761, 9910, -0: 12425, 13432, 12406, 10202, 11313, 12125, 12968, 13405, 12561, 11745, -0: 11194, 13006, 11227, 10395, 10354, 11110, 10107, 9996, 9598, 11741, -0: 13247, 11627, 8896, 9732, 12241, 11660, 10146, 10866, 10056, 9441, -0: 10619, 11767, 11657, 12170, 13797, 9380, 11389, 9277, 10227, 10911, -0: 9713, 10813, 10883, 10788, 10774, 10147, 8637, 9432, 10923, 9243, -0: 8978, 10170, 9043, 9322, 10205, 9454, 9322, 10353, 10208, 10225, -0: 8144, 10149, 9296, 9792, 9529, 8775, 7987, 9299, 7470, 8734, -0: 9667, 11008, 9733, 11530, 8974, 8481, 10398, 11173, 10894, 12060, -0: 8833, 9604, 7662, 9844, 10962, 9020, 9854, 9650, 10728, 9430, -0: 9986, 8931, 10908, 11575, 7910, 7259, 9472, 8202, 8697, 7909, -0: 7729, 7958, 7362, 8178, 8726, 9826, 8873, 7428, 8491, 7797, -0: 8774, 9180, 7704, 7044, 8096, 8939, 8507, 9518, 6994, 6797, -0: 9494, 8138, 8264, 7249, 8939, 7621, 7669, 8734, 7287, 7316, -0: 6354, 8748, 8245, 7315, 7517, 8989, 8748, 7412, 7965, 7501, -0: 7025, 7581, 7108, 9104, 8063, 9321, 8660, 8851, 7581, 9014, -0: 8373, 5826, 7538, 8144, 8096, 7110, 8038, 7664, 6487, 7958, -0: 7668, 9196, 9613, 6476, 8282, 7659, 10054, 6929, 7765, 7205, -0: 8737, 7037, 8079, 7974, 7485, 6498, 8750, 6987, 8448, 8182, -0: 8264, 7058, 7538, 6759, 7273, 7974, 7404, 7022, 6559, 8117, -0: 6684, 8929, 9311, 7385, 6911, 5942, 8717, 6315, 8271, 10950, -0: 8465, 6670, 8323, 8505, 7032, 6099, 5850, 5951, 5400, 5759, -0: 7351, 5581, 4874, 5359, 6755, 5394, 5361, 6433, 7439, 8102, -0: 7114, 7830, 6920, 6708, 6093, 5274, 7424, 5738, 5379, 7267, -0: 6562, 6738, 5992, 5474, 5726, 5788, 7554, 6488, 4801, 6940, -0: 6136, 6347, 5074, 5442, 6645, 5304, 5738, 5504, 5699, 5204, -0: 5812, 6853, 4638, 5094, 6053, 4599, 4470, 5836, 6226, 5780, -0: 5567, 5083, 7659, 5864, 5411, 7339, 5100, 4415, 5586, 5915, -0: 5691, 5845, 5445, 5372, 5219, 4534, 5999, 5901, 4948, 5483, -0: 5827, 5630, 5559, 7246, 5467, 4884, 5687, 5437, 6376, 5074, -0: 4927, 4594, 4565, 7177, 7211, 4619, 6537, 5087, 5818, 6429, -0: 6468, 5271, 5313, 6125, 6296, 5253, 5661, 5748, 4891, 6444, -0: 7318, 4465, 5150, 5627, 5165, 5245, 5743, 5536, 5678, 5381, -0: 5976, 6579, 4661, 5609, 4757, 5146, 5963, 5703, 5237, 5114, -0: 6305, 6249, 5101, 7171, 6257, 5286, 5902, 6002, 6414, 6107, -0: 6152, 7072, 6410, 5420, 6363, 5397, 5770, 5405, 5389, 6049, -0: 6478, 6551, 5583, 6499, 5315, 6228, 6329, 5793, 5633, 6642, -0: 5634, 5870, 5306, 4701, 4764, 5709, 5925, 5587, 5220, 6395, -0: 6056, 5686, 4966, 7269, 6085, 5358, 5729, 5793, 6671, 6677, -0: 5130, 5999, 4952, 4937, 4914, 6034, 6079, 4927, 5301, 5592, -0: 6145, 4639, 5447, 6324, 4967, 5862, 5804, 6168, 5382, 4937, -0: 5316, 6335, 5866, 6072, 4615, 5088, 6393, 5489, 5479, 5455, -0: 5791, 5293, 6251, 5061, 5485, 5937, 6079, 6906, 6176, 4499, -0: 4670, 4883, 4850, 4041, 4373, 4377, 4486, 4337, 4142, 4381, -0: 3568, 3777, 4208, 5078, 4130, 4806, 4168, 4886, 5110, 4768, -0: 3521, 4537, 4789, 3878, 4083, 4622, 4588, 4356, 4019, 3927, -0: 3691, 3659, 4114, 3882, 4676, 3806, 4460, 3033, 4273, 3014, -0: 4441, 4391, 4001, 4294, 4560, 4828, 4845, 5278, 4482, 4241, -0: 4978, 4077, 4034, 4583, 4253, 3464, 3590, 3784, 4414, 3643, -0: 3593, 3463, 3547, 5781, 3719, 4350, 3535, 3868, 3506, 3610, -0: 3783, 3316, 4171, 4885, 4319, 3630, 3866, 4652, 4751, 3500, -0: 3998, 3950], device='cuda:0', dtype=torch.int32) -0: construct-nextf: 0.785535991191864 -0: total_time: 1534.89404296875 -0: spgemm-matc-inst total_time: 0.0414079986512661 avg_time 0.0414079986512661 -0: spgemm-bcast-nnz total_time: 0.08691199868917465 avg_time 0.08691199868917465 -0: spgemm-inst-recv total_time: 3.5993599891662598 avg_time 3.5993599891662598 -0: spgemm-bcast-data total_time: 0.06860800087451935 avg_time 0.06860800087451935 -0: spgemm-preproc-local total_time: 0.22015999257564545 avg_time 0.22015999257564545 -0: spgemm-local-spgemm total_time: 96.85913848876953 avg_time 96.85913848876953 -0: spgemm-chunk-inst total_time: 1.1541119813919067 avg_time 1.1541119813919067 -0: spgemm-chunk-coalesce total_time: 33.36783981323242 avg_time 33.36783981323242 -0: spgemm-chunk-add total_time: 1.1110399961471558 avg_time 1.1110399961471558 -0: spgemm-matc-coalesce total_time: 0.01228800043463707 avg_time 0.01228800043463707 -0: spgemm-reduce-nnz total_time: 0.08790399879217148 avg_time 0.08790399879217148 -0: spgemm-padding total_time: 1.856160044670105 avg_time 1.856160044670105 -0: spgemm-allgather total_time: 2.243583917617798 avg_time 2.243583917617798 -0: spgemm-preproc-reduce total_time: 1.1478400230407715 avg_time 1.1478400230407715 -0: spgemm-reduce total_time: 39.09904098510742 avg_time 39.09904098510742 -0: spgemm-unpad total_time: 3.023871898651123 avg_time 3.023871898651123 -0: prob-rowsum total_time: 0.5768319964408875 avg_time 0.5768319964408875 -0: gen-darts total_time: 0.20255999267101288 avg_time 0.20255999267101288 -0: dart-throw total_time: 1.5697599649429321 avg_time 1.5697599649429321 -0: filter_darts total_time: 9.821632385253906 avg_time 9.821632385253906 -0: add-to-frontier total_time: 1.8423360586166382 avg_time 1.8423360586166382 -0: count_samples total_time: 30.70307159423828 avg_time 30.70307159423828 -0: select-preproc total_time: 0.07065600156784058 avg_time 0.07065600156784058 -0: select-psoverflow total_time: 0.6276479996740818 avg_time 0.06276479996740818 -0: select-reciprocal total_time: 13.165791749954224 avg_time 1.3165791749954223 -0: select-instmtx total_time: 11.41049611568451 avg_time 1.141049611568451 -0: select-invsum total_time: 591.7541084289551 avg_time 59.17541084289551 -0: select-psinv total_time: 0.6932480111718178 avg_time 0.06932480111718178 -0: select-computedarts total_time: 0.6707199960947037 avg_time 0.06707199960947037 -0: select-throwdarts total_time: 3.707904040813446 avg_time 0.3707904040813446 -0: select-add-to-frontier total_time: 16.77183973789215 avg_time 1.6771839737892151 -0: select-samplecount total_time: 310.00691413879395 avg_time 31.000691413879395 -0: select-overflow total_time: 0.695296011865139 avg_time 0.0695296011865139 -0: select-iter total_time: 951.7676544189453 avg_time 95.17676544189453 -0: dart-selection total_time: 957.3478393554688 avg_time 957.3478393554688 -0: set-probs total_time: 1.7666239738464355 avg_time 1.7666239738464355 -0: compute-bool total_time: 0.04710400104522705 avg_time 0.04710400104522705 -0: sampling-iters total_time: 1005.9694213867188 avg_time 1005.9694213867188 -0: select-mtxs total_time: 0.9820160269737244 avg_time 0.9820160269737244 -0: row-select-spgemm total_time: 197.51199340820312 avg_time 197.51199340820312 -0: row-select-expand total_time: 2.815295934677124 avg_time 2.815295934677124 -0: col-select-spgemm total_time: 79.22278594970703 avg_time 79.22278594970703 -0: set-sample total_time: 0.12403199821710587 avg_time 0.12403199821710587 -0: iter_count: 1 -0: selection_iter_count: 10 -0: -0: return_dart_count1.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], -0: device='cuda:0') -0: return_dart_count1: tensor([ 1, 1, 1, ..., 9, 5, 351], device='cuda:0', -0: dtype=torch.int32) -0: return_dart_map1: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: return_dart_val1: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, -0: 5.1157e+09], device='cuda:0') -0: return_dart_count2.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], -0: device='cuda:0') -0: return_dart_count2: tensor([ 1, 1, 1, ..., 9, 6, 304], device='cuda:0', -0: dtype=torch.int32) -0: return_dart_map2: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: return_dart_val2: tensor([3.9905e+06, 5.1668e+06, 2.4930e+05, ..., 5.1134e+09, 5.1199e+09, -0: 5.1157e+09], device='cuda:0') -0: -0: unmatched_darts: tensor([ 12, 43, 49, ..., 64487415, 64487416, 64487417], -0: device='cuda:0') -0: unmatched_darts.size: torch.Size([18767503]) -0: dart_map1[unmatched]: tensor([ 27788, 46136, 40383, ..., 36273593, 36267529, 36273527], -0: device='cuda:0', dtype=torch.int32) -0: dart_map2[unmatched]: tensor([ 27789, 46138, 40388, ..., 36273527, 36267509, 36273468], -0: device='cuda:0', dtype=torch.int32) -0: -0: dart_vals1[unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, -0: 5.1113e+09], device='cuda:0') -0: dart_vals2[unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, -0: 5.1113e+09], device='cuda:0') -0: -0: ps_pvals1[map1unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, -0: 5.1113e+09], device='cuda:0') -0: ps_pvals2[map1unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, -0: 5.1114e+09], device='cuda:0') -0: -0: ps_pvals1[map2unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1113e+09, 5.1103e+09, -0: 5.1113e+09], device='cuda:0') -0: ps_pvals2[map2unmatched]: tensor([4.0922e+06, 7.0450e+06, 5.6419e+06, ..., 5.1114e+09, 5.1103e+09, -0: 5.1113e+09], device='cuda:0') diff --git a/examples/tmp_bsearch_fpoint_error.out b/examples/tmp_bsearch_fpoint_error.out deleted file mode 100644 index 7cf3c2d..0000000 --- a/examples/tmp_bsearch_fpoint_error.out +++ /dev/null @@ -1,259 +0,0 @@ -0: Namespace(aggregator_type='gcn', batch_size=256, dataset='reddit', dist_backend='nccl', dist_url='env://', dropout=0.5, gpu=-1, hostname='127.0.0.1', lr=0.01, n_bulkmb=512, n_darts=-1, n_epochs=1, n_hidden=16, n_layers=1, normalize=False, partitioning='ONE5D', rank=-1, replication=1, samp_num=256, weight_decay=0.0, world_size=-1) -0: hostname: nid003637 -0: hostname: nid003637 rank: 0 size: 1 -0: start partitioning -0: rank: 0 adj_matrix_loc.size: torch.Size([232965, 232965]) -0: rank: 0 inputs_loc.size: torch.Size([232965, 602]) -0: end partitioning -0: coalescing -0: normalizing -0: done normalizing -0: imid: 18166615 val: 0.817233 arr[imid]: 256.478424 -0: imid: 9083307 val: 0.817233 arr[imid]: 128.845200 -0: imid: 4541653 val: 0.817233 arr[imid]: 64.056862 -0: imid: 2270826 val: 0.817233 arr[imid]: 31.714394 -0: imid: 1135412 val: 0.817233 arr[imid]: 15.978425 -0: imid: 567705 val: 0.817233 arr[imid]: 8.014808 -0: imid: 283852 val: 0.817233 arr[imid]: 4.157815 -0: imid: 141925 val: 0.817233 arr[imid]: 2.047342 -0: imid: 70962 val: 0.817233 arr[imid]: 0.983242 -0: imid: 35480 val: 0.817233 arr[imid]: 0.520355 -0: imid: 53221 val: 0.817233 arr[imid]: 0.788780 -0: imid: 62091 val: 0.817233 arr[imid]: 0.865206 -0: imid: 57656 val: 0.817233 arr[imid]: 0.833535 -0: imid: 55438 val: 0.817233 arr[imid]: 0.806440 -0: imid: 56547 val: 0.817233 arr[imid]: 0.811778 -0: imid: 57101 val: 0.817233 arr[imid]: 0.821501 -0: imid: 56824 val: 0.817233 arr[imid]: 0.817087 -0: imid: 56962 val: 0.817233 arr[imid]: 0.819000 -0: imid: 56893 val: 0.817233 arr[imid]: 0.817203 -0: imid: 56927 val: 0.817233 arr[imid]: 0.818839 -0: imid: 56910 val: 0.817233 arr[imid]: 0.817225 -0: imid: 56918 val: 0.817233 arr[imid]: 0.818824 -0: imid: 56914 val: 0.817233 arr[imid]: 0.817235 -0: imid: 56912 val: 0.817233 arr[imid]: 0.817234 -0: imid: 56911 val: 0.817233 arr[imid]: 0.817233 -0: imid: 18166615 val: 0.817233 arr[imid]: 256.478424 -0: imid: 9083307 val: 0.817233 arr[imid]: 128.845184 -0: imid: 4541653 val: 0.817233 arr[imid]: 64.056938 -0: imid: 2270826 val: 0.817233 arr[imid]: 31.714397 -0: imid: 1135412 val: 0.817233 arr[imid]: 15.978395 -0: imid: 567705 val: 0.817233 arr[imid]: 8.014786 -0: imid: 283852 val: 0.817233 arr[imid]: 4.157809 -0: imid: 141925 val: 0.817233 arr[imid]: 2.047338 -0: imid: 70962 val: 0.817233 arr[imid]: 0.983241 -0: imid: 35480 val: 0.817233 arr[imid]: 0.520354 -0: imid: 53221 val: 0.817233 arr[imid]: 0.788780 -0: imid: 62091 val: 0.817233 arr[imid]: 0.865205 -0: imid: 57656 val: 0.817233 arr[imid]: 0.833534 -0: imid: 55438 val: 0.817233 arr[imid]: 0.806439 -0: imid: 56547 val: 0.817233 arr[imid]: 0.811777 -0: imid: 57101 val: 0.817233 arr[imid]: 0.821501 -0: imid: 56824 val: 0.817233 arr[imid]: 0.817087 -0: imid: 56962 val: 0.817233 arr[imid]: 0.818999 -0: imid: 56893 val: 0.817233 arr[imid]: 0.817202 -0: imid: 56927 val: 0.817233 arr[imid]: 0.818838 -0: imid: 56910 val: 0.817233 arr[imid]: 0.817224 -0: imid: 56918 val: 0.817233 arr[imid]: 0.818824 -0: imid: 56914 val: 0.817233 arr[imid]: 0.817234 -0: imid: 56912 val: 0.817233 arr[imid]: 0.817233 -0: imid: 56913 val: 0.817233 arr[imid]: 0.817233 -0: n_darts: 125952 -0: rank: 0 g_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 232964, 232964, 232964], -0: [ 0, 242, 249, ..., 232594, 232634, 232964]]), -0: values=tensor([2.0568e-07, 2.0568e-07, 2.0568e-07, ..., -0: 4.1311e-06, 4.1311e-06, 4.1311e-06]), -0: device='cuda:0', size=(232965, 232965), nnz=114848857, -0: layout=torch.sparse_coo) -0: rank: 0 batches_loc: tensor(indices=tensor([[ 0, 0, 0, ..., 511, 511, 511], -0: [195991, 117560, 100852, ..., 210813, 53334, 131122]]), -0: values=tensor([1., 1., 1., ..., 1., 1., 1.]), -0: device='cuda:0', size=(512, 232965), nnz=131072, layout=torch.sparse_coo) -0: probability-spgemm: 352.7101745605469 -0: compute-p: 60.275230407714844 -0: p.nnz: 36333230 -0: pre-loop: 1.4532480239868164 -0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') -0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], -0: [ 8], -0: [ 11], -0: ..., -0: [36332465], -0: [36332476], -0: [36332500]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3814725, 1]) -0: iter_count: 1 dart_hits_count.sum: 64487424 -0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[ 1], -0: [ 1], -0: [ 1], -0: ..., -0: [ 14], -0: [ 3], -0: [274]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9509e-08, 2.2447e-08, ..., 5.1202e+02, 5.1202e+02, -0: 5.1202e+02], device='cuda:0') -0: construct-nextf: 0.7938240170478821 -0: total_time: 1601.1064453125 -0: spgemm-matc-inst total_time: 0.059647999703884125 avg_time 0.059647999703884125 -0: spgemm-bcast-nnz total_time: 132.69155883789062 avg_time 132.69155883789062 -0: spgemm-inst-recv total_time: 3.57206392288208 avg_time 3.57206392288208 -0: spgemm-bcast-data total_time: 0.08755200356245041 avg_time 0.08755200356245041 -0: spgemm-preproc-local total_time: 0.35334399342536926 avg_time 0.35334399342536926 -0: spgemm-local-spgemm total_time: 96.30374145507812 avg_time 96.30374145507812 -0: spgemm-chunk-inst total_time: 1.1650559902191162 avg_time 1.1650559902191162 -0: spgemm-chunk-coalesce total_time: 33.32175827026367 avg_time 33.32175827026367 -0: spgemm-chunk-add total_time: 1.1077120304107666 avg_time 1.1077120304107666 -0: spgemm-matc-coalesce total_time: 0.012768000364303589 avg_time 0.012768000364303589 -0: spgemm-reduce-nnz total_time: 36.258689880371094 avg_time 36.258689880371094 -0: spgemm-padding total_time: 1.8590400218963623 avg_time 1.8590400218963623 -0: spgemm-allgather total_time: 2.2658560276031494 avg_time 2.2658560276031494 -0: spgemm-preproc-reduce total_time: 1.1514240503311157 avg_time 1.1514240503311157 -0: spgemm-reduce total_time: 39.014591217041016 avg_time 39.014591217041016 -0: spgemm-unpad total_time: 3.0386879444122314 avg_time 3.0386879444122314 -0: prob-rowsum total_time: 0.5860480070114136 avg_time 0.5860480070114136 -0: gen-darts total_time: 0.2377600073814392 avg_time 0.2377600073814392 -0: dart-throw total_time: 1.3964799642562866 avg_time 1.3964799642562866 -0: filter_darts total_time: 9.814399719238281 avg_time 9.814399719238281 -0: add-to-frontier total_time: 1.8496639728546143 avg_time 1.8496639728546143 -0: count_samples total_time: 30.425567626953125 avg_time 30.425567626953125 -0: select-preproc total_time: 0.14537599682807922 avg_time 0.14537599682807922 -0: select-psoverflow total_time: 0.5724479965865612 avg_time 0.06360533295406236 -0: select-reciprocal total_time: 11.853407859802246 avg_time 1.317045317755805 -0: select-instmtx total_time: 10.271583795547485 avg_time 1.141287088394165 -0: select-invsum total_time: 531.1580810546875 avg_time 59.01756456163194 -0: select-psinv total_time: 0.6277760118246078 avg_time 0.0697528902027342 -0: select-computedarts total_time: 0.6289280094206333 avg_time 0.06988088993562593 -0: select-throwdarts total_time: 3.477055996656418 avg_time 0.38633955518404645 -0: select-add-to-frontier total_time: 15.113856077194214 avg_time 1.679317341910468 -0: select-samplecount total_time: 277.2289619445801 avg_time 30.80321799384223 -0: select-overflow total_time: 0.6320320069789886 avg_time 0.07022577855322096 -0: select-iter total_time: 853.6518630981445 avg_time 94.85020701090495 -0: dart-selection total_time: 853.9918212890625 avg_time 853.9918212890625 -0: set-probs total_time: 1.7295680046081543 avg_time 1.7295680046081543 -0: compute-bool total_time: 0.055904000997543335 avg_time 0.055904000997543335 -0: sampling-iters total_time: 902.2776489257812 avg_time 902.2776489257812 -0: select-mtxs total_time: 0.9269440174102783 avg_time 0.9269440174102783 -0: row-select-spgemm total_time: 200.27154541015625 avg_time 200.27154541015625 -0: row-select-expand total_time: 2.836479902267456 avg_time 2.836479902267456 -0: col-select-spgemm total_time: 79.11392211914062 avg_time 79.11392211914062 -0: set-sample total_time: 0.131071999669075 avg_time 0.131071999669075 -0: iter_count: 1 -0: selection_iter_count: 9 -0: -0: probability-spgemm: 184.49417114257812 -0: compute-p: 60.3001594543457 -0: p.nnz: 36333230 -0: pre-loop: 1.4463039636611938 -0: iter_count: 1 dart_values: tensor([0.3990, 0.5167, 0.0249, ..., 0.3437, 0.9927, 0.5703], device='cuda:0') -0: iter_count: 1 dart_values_scaled: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs: tensor([[ 4], -0: [ 8], -0: [ 11], -0: ..., -0: [36332465], -0: [36332476], -0: [36332500]], device='cuda:0') -0: iter_count: 1 dart_hits_count.nnzidxs.size: torch.Size([3813474, 1]) -0: iter_count: 1 dart_hits_count.sum: 64487424 -0: iter_count: 1 dart_hits_map: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: iter_count: 1 dart_hits_count: tensor([[ 1], -0: [ 1], -0: [ 1], -0: ..., -0: [ 12], -0: [ 3], -0: [271]], device='cuda:0', dtype=torch.int32) -0: iter_count: 1 ps_p_values: tensor([0.0000e+00, 1.9509e-08, 2.2447e-08, ..., 5.1202e+02, 5.1202e+02, -0: 5.1202e+02], device='cuda:0') -0: construct-nextf: 0.8681280016899109 -0: total_time: 1430.2353515625 -0: spgemm-matc-inst total_time: 0.043327998369932175 avg_time 0.043327998369932175 -0: spgemm-bcast-nnz total_time: 0.0880960002541542 avg_time 0.0880960002541542 -0: spgemm-inst-recv total_time: 3.5996479988098145 avg_time 3.5996479988098145 -0: spgemm-bcast-data total_time: 0.06771200150251389 avg_time 0.06771200150251389 -0: spgemm-preproc-local total_time: 0.2189120054244995 avg_time 0.2189120054244995 -0: spgemm-local-spgemm total_time: 96.97039794921875 avg_time 96.97039794921875 -0: spgemm-chunk-inst total_time: 1.1507200002670288 avg_time 1.1507200002670288 -0: spgemm-chunk-coalesce total_time: 33.33875274658203 avg_time 33.33875274658203 -0: spgemm-chunk-add total_time: 1.1117119789123535 avg_time 1.1117119789123535 -0: spgemm-matc-coalesce total_time: 0.01190400030463934 avg_time 0.01190400030463934 -0: spgemm-reduce-nnz total_time: 0.09206400066614151 avg_time 0.09206400066614151 -0: spgemm-padding total_time: 1.8569920063018799 avg_time 1.8569920063018799 -0: spgemm-allgather total_time: 2.242719888687134 avg_time 2.242719888687134 -0: spgemm-preproc-reduce total_time: 1.1456639766693115 avg_time 1.1456639766693115 -0: spgemm-reduce total_time: 39.106143951416016 avg_time 39.106143951416016 -0: spgemm-unpad total_time: 3.037760019302368 avg_time 3.037760019302368 -0: prob-rowsum total_time: 0.5271360278129578 avg_time 0.5271360278129578 -0: gen-darts total_time: 0.20892800390720367 avg_time 0.20892800390720367 -0: dart-throw total_time: 1.305791974067688 avg_time 1.305791974067688 -0: filter_darts total_time: 9.789440155029297 avg_time 9.789440155029297 -0: add-to-frontier total_time: 1.8480960130691528 avg_time 1.8480960130691528 -0: count_samples total_time: 30.420608520507812 avg_time 30.420608520507812 -0: select-preproc total_time: 0.07011199742555618 avg_time 0.07011199742555618 -0: select-psoverflow total_time: 0.5702720023691654 avg_time 0.06336355581879616 -0: select-reciprocal total_time: 11.712607979774475 avg_time 1.3014008866416082 -0: select-instmtx total_time: 10.291007995605469 avg_time 1.143445332845052 -0: select-invsum total_time: 530.8290901184082 avg_time 58.981010013156464 -0: select-psinv total_time: 0.648159995675087 avg_time 0.07201777729723188 -0: select-computedarts total_time: 0.6293440014123917 avg_time 0.06992711126804352 -0: select-throwdarts total_time: 3.543071985244751 avg_time 0.39367466502719456 -0: select-add-to-frontier total_time: 15.13158392906189 avg_time 1.6812871032290988 -0: select-samplecount total_time: 277.18652725219727 avg_time 30.79850302802192 -0: select-overflow total_time: 0.641087993979454 avg_time 0.07123199933105046 -0: select-iter total_time: 853.2802658081055 avg_time 94.80891842312283 -0: dart-selection total_time: 853.6178588867188 avg_time 853.6178588867188 -0: set-probs total_time: 1.7106560468673706 avg_time 1.7106560468673706 -0: compute-bool total_time: 0.045343998819589615 avg_time 0.045343998819589615 -0: sampling-iters total_time: 901.5848999023438 avg_time 901.5848999023438 -0: select-mtxs total_time: 0.9415680170059204 avg_time 0.9415680170059204 -0: row-select-spgemm total_time: 198.17837524414062 avg_time 198.17837524414062 -0: row-select-expand total_time: 2.817023992538452 avg_time 2.817023992538452 -0: col-select-spgemm total_time: 79.16553497314453 avg_time 79.16553497314453 -0: set-sample total_time: 0.13356800377368927 avg_time 0.13356800377368927 -0: iter_count: 1 -0: selection_iter_count: 9 -0: -0: return_dart_count1.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], -0: device='cuda:0') -0: return_dart_count1: tensor([ 1, 1, 1, ..., 14, 3, 274], device='cuda:0', -0: dtype=torch.int32) -0: return_dart_map1: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: return_dart_val1: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: return_dart_count2.nnzidxs: tensor([ 4, 8, 11, ..., 36332465, 36332476, 36332500], -0: device='cuda:0') -0: return_dart_count2: tensor([ 1, 1, 1, ..., 12, 3, 271], device='cuda:0', -0: dtype=torch.int32) -0: return_dart_map2: tensor([ 27406, 34971, 3442, ..., 36296107, 36331806, 36309730], -0: device='cuda:0', dtype=torch.int32) -0: return_dart_val2: tensor([3.9905e-01, 5.1668e-01, 2.4930e-02, ..., 5.1134e+02, 5.1199e+02, -0: 5.1157e+02], device='cuda:0') -0: -0: unmatched_darts: tensor([ 92, 93, 106, ..., 64487413, 64487416, 64487417], -0: device='cuda:0') -0: dart_map1[unmatched]: tensor([ 56912, 45262, 50443, ..., 36330476, 36267481, 36273438], -0: device='cuda:0', dtype=torch.int32) -0: dart_map2[unmatched]: tensor([ 56914, 45263, 50444, ..., 36330467, 36267509, 36273436], -0: device='cuda:0', dtype=torch.int32) -0: -0: dart_vals1[unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9604, 511.0260, 511.1345], -0: device='cuda:0') -0: dart_vals2[unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9604, 511.0260, 511.1345], -0: device='cuda:0') -0: -0: ps_pvals1[map1unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9605, 511.0261, 511.1346], -0: device='cuda:0') -0: ps_pvals2[map1unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9606, 511.0259, 511.1347], -0: device='cuda:0') -0: -0: ps_pvals1[map2unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9604, 511.0264, 511.1344], -0: device='cuda:0') -0: ps_pvals2[map2unmatched]: tensor([ 0.8172, 0.6601, 0.7475, ..., 511.9605, 511.0262, 511.1346], -0: device='cuda:0') From 35a15a43b2041fa8b62974796b02ada004cc0e3f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 9 Aug 2022 23:25:01 -0700 Subject: [PATCH 083/178] Use avg_deg * batch_size as dart count --- examples/gcn_15d.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index ddc2761..427840e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -369,19 +369,6 @@ def main(args): edge_count = adj_matrix.size(1) avg_degree = int(edge_count / node_count) args.n_darts = avg_degree * args.batch_size - - degs = torch.cuda.IntTensor(node_count).fill_(0) - degs.scatter_add_(0, adj_matrix[0, :].long(), torch.cuda.IntTensor(edge_count).fill_(1)) - degs = degs[degs.nonzero()] - print(f"degs: {degs.squeeze()}") - max_degree = torch.max(degs).item() - med_degree = torch.median(degs).item() - print(f"degs.size: {degs.size()}") - print(f"degs.nnz.size: {degs.nonzero().squeeze().size()}") - - print(f"edge_count: {edge_count} node_count: {node_count} avg_degree: {avg_degree} max_degree: {max_degree} med_degree: {med_degree}") - - args.n_darts = int(max_degree / avg_degree) print(f"n_darts: {args.n_darts}") print(f"rank: {rank} g_loc: {g_loc}") From 6ae29bc8750b5dc70aad1fafa8f409e4867c5145 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 10 Aug 2022 13:23:01 -0700 Subject: [PATCH 084/178] Updated to torch 1.11 for nccl gather() --- cagnet/samplers/ladies.py | 1 + examples/__init__.py | 0 sparse-extension/sparse_coo_tensor.cu | 4 +- tests/test_nccl_p2p.py | 55 +++++++++++++++++++++++---- 4 files changed, 50 insertions(+), 10 deletions(-) create mode 100644 examples/__init__.py diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 8e38dc9..2da2a2a 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -49,6 +49,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 03480b1..0023c3b 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -11,7 +11,7 @@ #include -#include +// #include #include @@ -173,7 +173,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, // cusparseHandle_t handle; // CHECK_CUSPARSE(cusparseCreate(&handle)); - auto state = at::globalContext().lazyInitCUDA(); + // auto state = at::globalContext().lazyInitCUDA(); // auto handle = THCState_getCurrentSparseHandle(state); auto handle = at::cuda::getCurrentCUDASparseHandle(); diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index a6030ae..059d947 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -20,14 +20,53 @@ def test_nccl(args): print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}") print(f"backend: {dist.get_backend()}") - if rank == 0: - x = torch.cuda.IntTensor(5).fill_(3) # vector of size 5, set to all 3's - dist.send(x, dst=1) - elif rank == 1: - y = torch.cuda.IntTensor(5).fill_(2) # vector of size 5, set to all 2's - print(f"y_before: {y}") - dist.recv(y, src=0) - print(f"y_after: {y}") + # # send/recv test + # if rank == 0: + # x = torch.cuda.IntTensor(5).fill_(3) # vector of size 5, set to all 3's + # dist.send(x, dst=1) + # elif rank == 1: + # y = torch.cuda.IntTensor(5).fill_(2) # vector of size 5, set to all 2's + # print(f"y_before: {y}") + # dist.recv(y, src=0) + # print(f"y_after: {y}") + + # # isend/irecv test + # if rank == 0: + # x = torch.cuda.IntTensor(rank + 1).fill_(0) # vector of size 5, set to all 3's + # recv_objs = [] + # recv_tens = [] + # for i in range(size): + # recv_tens.append(torch.cuda.IntTensor(i + 1).fill_(0)) + + # print(f"before recv_tens: {recv_tens}") + # for i in range(1, size): + # print(f"i: {i}") + # recv_objs.append(dist.irecv(recv_tens[i], src=i)) + # # map(lambda x:x.wait(), recv_objs) + # for obj in recv_objs: + # obj.wait() + # print(f"after recv_tens: {recv_tens}") + # else: + # y = torch.cuda.IntTensor(rank + 1).fill_(rank) # vector of size 5, set to all 2's + # print(f"rank: {rank} y: {y}") + # send_obj = dist.isend(y, dst=0) + # send_obj.wait() + + # gather test + if rank > 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}") + if __name__ == '__main__': parser = argparse.ArgumentParser(description='NCCL P2P test') From d023bb3d08664f3a3a6e2d4bf96ca0b7c1b757b0 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 11 Aug 2022 18:42:15 -0700 Subject: [PATCH 085/178] Sparsity-aware 1.5D matches 1.5D --- .gitignore | 1 + cagnet/samplers/ladies.py | 180 +++++++++++++++++++++++++- sparse-extension/sparse_coo_tensor.cu | 52 +++++++- 3 files changed, 231 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 932b594..4bf800c 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,4 @@ tests/examples tests/LADIES/ tests/*.out *.qdrep +*.out diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 2da2a2a..5c8a1c2 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -8,7 +8,7 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu + scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu def start_time(timer): timer.record() @@ -149,6 +149,175 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, return matc_nnz_indices, matc_nnz_values +def dist_saspgemm15D(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}") + + 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 + 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}") + + # Determine number of nonzero columns in chunk of mata + nnz_cols = torch.unique(mata_chunk_indices[1,:]) + nnz_cols_count = torch.cuda.IntTensor(1).fill_(nnz_cols.size(0)) + + # Gather nnz column counts on rank q + if rank == q: + nnz_cols_count_list = [] + for j in range(size // replication): + nnz_cols_count_list.append(torch.cuda.IntTensor(1).fill_(0)) + 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]) + + # Rank q allocates recv buffer for nnz col ids + 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)) + + # isend/irecv nnz col ids to rank q + 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.irecv(nnz_col_ids[j], src=recv_rank)) + else: + nnz_col_ids[j] = nnz_cols + for obj in recv_objs: + obj.wait() + else: + dist.isend(nnz_cols, dst=q).wait() + + # Send selected rows count selected rows from block row of matb from nnz_col_ids + if rank == q: + send_objs = [] + for j in range(size // replication): + recv_rank = rank_col + j * replication + nnz_row_mask = torch.cuda.LongTensor(matb._indices().size(1)).fill_(0) + # nnz_row_mask[nnz_col_ids[j].long()] = True + rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ + nnz_col_ids[j].size(0), matb._indices().size(1)) + nnz_row_mask = nnz_row_mask.bool() + + matb_send_indices = matb._indices()[:, nnz_row_mask].float() + matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) + matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) + + if recv_rank != q: + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) + send_objs.append(dist.isend(selected_rows_count, tag=0, dst=recv_rank)) + send_objs.append(dist.isend(matb_send, tag=1, dst=recv_rank)) + else: + matb_select_recv = matb_send.clone() + + for obj in send_objs: + obj.wait() + else: + selected_rows_count_recv = torch.cuda.IntTensor(1).fill_(0) + dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() + + matb_select_recv = torch.cuda.DoubleTensor(3, selected_rows_count_recv.item()).fill_(0) + dist.irecv(matb_select_recv, tag=1, src=q).wait() + + matb_recv_indices = matb_select_recv[:2, :].long() + matb_recv_values = matb_select_recv[2, :].double() + + start_time(start_timer) + 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) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + + start_time(start_timer) + matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + + start_time(start_timer) + matc_chunk = matc_chunk.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + + start_time(start_timer) + matc += matc_chunk + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") + + # 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}") + + 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}") + + 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}") + + 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())) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}") + + 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}") + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") + + start_time(start_timer) + nnz_mask = matc._values() != -1.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}") + + return matc_nnz_indices, matc_nnz_values + def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ replication, rank, size, row_groups, col_groups): @@ -193,6 +362,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-probability-spgemm") p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, \ row_groups, col_groups, "prob", timing_dict) + # p_num_indices, p_num_values = dist_saspgemm15D(batches, adj_matrix, replication, rank, size, \ + # row_groups, col_groups, "prob", timing_dict) torch.cuda.nvtx.range_pop() timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -206,6 +377,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota values=p_num_values, size=(mb_count, node_count_total)) print(f"p.nnz: {p._nnz()}") + print(f"p.indices: {p._indices()}") + print(f"p.indices[row 0].size: {(p._indices()[0,:] == 0).nonzero().squeeze().size()}") + print(f"p.indices[row 1].size: {(p._indices()[0,:] == 1).nonzero().squeeze().size()}") + print(f"p.indices[row 2].size: {(p._indices()[0,:] == 2).nonzero().squeeze().size()}") + print(f"p.indices[row 3].size: {(p._indices()[0,:] == 3).nonzero().squeeze().size()}") + print(f"p.values: {p._values()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) @@ -331,6 +508,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # Compute darts for selection compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) + print(flush=True) timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 0023c3b..7bbad31 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -74,7 +74,7 @@ __device__ long binary_searchl(long *arr, long val, long imin, long imax) { if (arr[imid] <= val) { imin = imid + 1; - } else { + } else if (arr[imid] > val) { ans = imid; imax = imid - 1; } @@ -83,6 +83,23 @@ __device__ long binary_searchl(long *arr, long val, long imin, long imax) { 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) { @@ -658,6 +675,38 @@ void scatteri_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at CHECK_ERROR("scatter add ints error") } +__global__ void RowSelectCoo(long *nnz_cols, long *row_ids, long *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] = 1L; + } + } +} + +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(nnz_col_count / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (nnz_col_count == 0) { + return; + } + + RowSelectCoo<<>>(nnz_cols.data(), + rows.data(), + mask.data(), + nnz_col_count, + row_count); + CHECK_ERROR("rowselect coo error") +} + 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"); @@ -673,4 +722,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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, "Row selection for sparsity-aware spgemm"); } From a44118cd52555afb90b9697725968722cb9a07b2 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 11 Aug 2022 19:20:10 -0700 Subject: [PATCH 086/178] Passes test script --- cagnet/samplers/ladies.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 5c8a1c2..b3d4cd6 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -142,7 +142,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") 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}") @@ -311,7 +311,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") start_time(start_timer) - nnz_mask = matc._values() != -1.0 + # 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}") From 43d63b29694aa4f980094514d12ceff7c312427b Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 11 Aug 2022 19:20:50 -0700 Subject: [PATCH 087/178] Cleanup --- cagnet/samplers/ladies.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index b3d4cd6..ff2bebe 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -378,12 +378,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota values=p_num_values, size=(mb_count, node_count_total)) print(f"p.nnz: {p._nnz()}") - print(f"p.indices: {p._indices()}") - print(f"p.indices[row 0].size: {(p._indices()[0,:] == 0).nonzero().squeeze().size()}") - print(f"p.indices[row 1].size: {(p._indices()[0,:] == 1).nonzero().squeeze().size()}") - print(f"p.indices[row 2].size: {(p._indices()[0,:] == 2).nonzero().squeeze().size()}") - print(f"p.indices[row 3].size: {(p._indices()[0,:] == 3).nonzero().squeeze().size()}") - print(f"p.values: {p._values()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) From 7c02c5dbda99150768e2865f49f3009f895c6925 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 11 Aug 2022 19:22:09 -0700 Subject: [PATCH 088/178] Cleanup --- cagnet/samplers/ladies.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index ff2bebe..0458901 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -503,7 +503,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # Compute darts for selection compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, mb_count, total_overflow) - print(flush=True) timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) From f75908da1b457f740208d440bc0986d051370d58 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 12 Aug 2022 12:47:13 -0700 Subject: [PATCH 089/178] Correct but slow sparsity-aware 1.5D spgemm --- .gitignore | 1 + cagnet/samplers/ladies.py | 135 ++++++++++++++++++++----- examples/gcn_15d.py | 13 ++- sparse-extension/sparse_coo_tensor.cu | 6 +- tests/test_nccl.cu | 139 ++++++++++++++++++++++++++ tests/test_nccl_p2p.py | 96 +++++++++++------- 6 files changed, 321 insertions(+), 69 deletions(-) create mode 100644 tests/test_nccl.cu diff --git a/.gitignore b/.gitignore index 4bf800c..15f95b3 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,4 @@ tests/LADIES/ tests/*.out *.qdrep *.out +*.nsys-rep diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 0458901..414dc94 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -20,19 +20,27 @@ def stop_time(start_timer, stop_timer): def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if timing_dict is not None: - misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ - not range_name.startswith("spgemm-bcast-data") and \ - not range_name.startswith("spgemm-local-spgemm") and \ - not range_name.startswith("spgemm-reduce-nnz") and \ - not range_name.startswith("spgemm-reduce") - if misc_range: + # misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ + # not range_name.startswith("spgemm-bcast-data") and \ + # not range_name.startswith("spgemm-local-spgemm") and \ + # not range_name.startswith("spgemm-reduce-nnz") and \ + # not range_name.startswith("spgemm-reduce") + if False and misc_range: timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) else: timing_dict[range_name].append(stop_time(start_timer, stop_timer)) +saspgemm_data = 0 +spgemm_data = 0 + def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ name, timing_dict=None): + global spgemm_data + + if name == "prob": + spgemm_data = 0 + start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -68,6 +76,9 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + if name == "prob": + spgemm_data += matb_recv_indices.numel() + matb_recv_values.numel() + start_time(start_timer) am_partid = rank_col * (size // replication ** 2) + i chunk_col_start = am_partid * chunk_size @@ -150,11 +161,17 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, return matc_nnz_indices, matc_nnz_values def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ - name, timing_dict=None): + name, nnz_row_masks, matb_recv_buff, timing_dict=None): + + global saspgemm_data + saspgemm_data = 0 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))) matc = torch.sparse_coo_tensor(size=(mata.size(0), matb.size(1))).cuda() @@ -183,10 +200,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}") # Determine number of nonzero columns in chunk of mata + start_time(start_timer) nnz_cols = torch.unique(mata_chunk_indices[1,:]) 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}") # Gather nnz column counts on rank q + start_time(start_timer) if rank == q: nnz_cols_count_list = [] for j in range(size // replication): @@ -194,14 +214,18 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") # 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}") # isend/irecv nnz col ids to rank q + start_time(start_timer) if rank == q: recv_objs = [] for j in range(size // replication): @@ -214,37 +238,87 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups obj.wait() else: dist.isend(nnz_cols, dst=q).wait() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") - # Send selected rows count selected rows from block row of matb from nnz_col_ids + # Send selected rows count + torch.cuda.nvtx.range_push("nvtx-send-nnzrows") + start_time(start_timer) if rank == q: - send_objs = [] + send_objs = [None] * (size // replication) + matb_send_list = [None] * (size // replication) for j in range(size // replication): recv_rank = rank_col + j * replication - nnz_row_mask = torch.cuda.LongTensor(matb._indices().size(1)).fill_(0) - # nnz_row_mask[nnz_col_ids[j].long()] = True + start_time(start_inner_timer) + nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") + start_time(start_inner_timer) rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ nnz_col_ids[j].size(0), matb._indices().size(1)) - nnz_row_mask = nnz_row_mask.bool() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") - matb_send_indices = matb._indices()[:, nnz_row_mask].float() + start_time(start_inner_timer) + matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskindalc-{name}") + + start_time(start_inner_timer) + matb_send_indices = matb._indices()[:, nnz_row_mask] + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskind-{name}") + + start_time(start_inner_timer) + matb_send_indices = matb_send_indices.float() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-indfloat-{name}") + + start_time(start_inner_timer) matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskval-{name}") + + start_time(start_inner_timer) matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") + + matb_send_list[j] = matb_send if recv_rank != q: selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) - send_objs.append(dist.isend(selected_rows_count, tag=0, dst=recv_rank)) - send_objs.append(dist.isend(matb_send, tag=1, dst=recv_rank)) + send_objs[j] = dist.isend(selected_rows_count, dst=recv_rank) else: matb_select_recv = matb_send.clone() + start_time(start_inner_timer) for obj in send_objs: - obj.wait() + if obj is not None: + obj.wait() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-wait-{name}") else: - selected_rows_count_recv = torch.cuda.IntTensor(1).fill_(0) - dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() + start_time(start_inner_timer) + selected_rows_count_recv = torch.cuda.IntTensor(1) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") + dist.irecv(selected_rows_count_recv, src=q).wait() - matb_select_recv = torch.cuda.DoubleTensor(3, selected_rows_count_recv.item()).fill_(0) - dist.irecv(matb_select_recv, tag=1, src=q).wait() + start_time(start_inner_timer) + matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") + + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") + torch.cuda.nvtx.range_pop() + + # Send selected rows from block row of matb from nnz_col_ids + start_time(start_timer) + if rank == q: + send_objs = [None] * (size // replication) + for j in range(size // replication): + recv_rank = rank_col + j * replication + if recv_rank != q: + send_objs[j] = dist.isend(matb_send_list[j], dst=recv_rank) + if name == "prob": + saspgemm_data += matb_send.numel() + for obj in send_objs: + if obj is not None: + obj.wait() + else: + dist.irecv(matb_select_recv, src=q).wait() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzrows-{name}") + torch.cuda.nvtx.range_pop() matb_recv_indices = matb_select_recv[:2, :].long() matb_recv_values = matb_select_recv[2, :].double() @@ -320,7 +394,10 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups return matc_nnz_indices, matc_nnz_values def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ - replication, rank, size, row_groups, col_groups): + replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups): + + global spgemm_data + global saspgemm_data start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -361,10 +438,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, \ - row_groups, col_groups, "prob", timing_dict) - # p_num_indices, p_num_values = dist_saspgemm15D(batches, adj_matrix, replication, rank, size, \ + # p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, \ # row_groups, col_groups, "prob", timing_dict) + p_num_indices, p_num_values = dist_saspgemm15D(batches, adj_matrix, replication, rank, size, \ + row_groups, col_groups, "prob", sa_masks, sa_recv_buff, + timing_dict) torch.cuda.nvtx.range_pop() timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -651,13 +729,13 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_pop() print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) - for k, v in timing_dict.items(): + 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": # print(f"{k} times: {v}") - v_tens = torch.cuda.FloatTensor([sum(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)) + v_tens_recv.append(torch.cuda.FloatTensor(1).fill_(0)) dist.all_gather(v_tens_recv, v_tens) if rank == 0: @@ -667,8 +745,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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(): print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)} len: {len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") + print(f"spgemm_data: {spgemm_data}") + print(f"saspgemm_data: {saspgemm_data}") return batches_select, next_frontier_select, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 427840e..69a40a1 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -375,12 +375,17 @@ def main(args): print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing torch.manual_seed(0) + nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm + nnz_row_masks.fill_(0) + + nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) + sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) current_frontier, next_frontier, adj_matrices = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, rank, size, \ - row_groups, col_groups) + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() @@ -390,8 +395,8 @@ def main(args): ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, rank, size, \ - row_groups, col_groups) + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 7bbad31..8e6e8b5 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -675,14 +675,14 @@ void scatteri_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at CHECK_ERROR("scatter add ints error") } -__global__ void RowSelectCoo(long *nnz_cols, long *row_ids, long *mask, int nnz_col_count, int row_count) { +__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] = 1L; + mask[i] = true; } } } @@ -701,7 +701,7 @@ void rowselect_coo_gpu(const at::Tensor& nnz_cols, const at::Tensor& rows, const RowSelectCoo<<>>(nnz_cols.data(), rows.data(), - mask.data(), + mask.data(), nnz_col_count, row_count); CHECK_ERROR("rowselect coo error") diff --git a/tests/test_nccl.cu b/tests/test_nccl.cu new file mode 100644 index 0000000..fb9faa4 --- /dev/null +++ b/tests/test_nccl.cu @@ -0,0 +1,139 @@ +#include +#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) + # isend/irecv test + # buffer_size = 344546571 # spgemm + # buffer_size = 37135566 # sa-spgemm + buffer_size = 37135566 // 4 # sa-spgemm single isend + # Communicate once w/o timing + if rank == 0: + recv_objs = [] + recv_tens = torch.cuda.LongTensor(buffer_size).fill_(2) + + print(f"before recv_tens: {recv_tens}") + # dist.irecv(recv_tens, src=1).wait() + dist.recv(recv_tens, src=1) + print(f"after recv_tens: {recv_tens}") 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}") + y = torch.cuda.LongTensor(buffer_size).fill_(3) # vector of size 5, set to all 2's + # dist.isend(y, dst=0).wait() + dist.send(y, dst=0) + + if rank == 0: + recv_objs = [] + recv_tens = torch.cuda.LongTensor(buffer_size).fill_(5) + + print(f"before recv_tens: {recv_tens}") + # dist.irecv(recv_tens, src=1).wait() + dist.recv(recv_tens, src=1) + print(f"after recv_tens: {recv_tens}") + else: + y = torch.cuda.LongTensor(buffer_size).fill_(6) # vector of size 5, set to all 2's + start_time(start_timer) + # dist.isend(y, dst=0).wait() + dist.send(y, dst=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}") + + # # gather test + # if rank > 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}") if __name__ == '__main__': From 7b67556be9bfd0f4e8a8dc02bb80fc3ed439818f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 25 Aug 2022 15:21:35 -0700 Subject: [PATCH 090/178] Performance debugging SA-SpGEMM --- cagnet/samplers/ladies.py | 72 +++++++++++---------------- sparse-extension/sparse_coo_tensor.cu | 2 +- 2 files changed, 30 insertions(+), 44 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 414dc94..12d3923 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -234,8 +234,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups recv_objs.append(dist.irecv(nnz_col_ids[j], src=recv_rank)) else: nnz_col_ids[j] = nnz_cols - for obj in recv_objs: - obj.wait() else: dist.isend(nnz_cols, dst=q).wait() stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") @@ -244,80 +242,68 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-send-nnzrows") start_time(start_timer) if rank == q: - send_objs = [None] * (size // replication) - matb_send_list = [None] * (size // replication) for j in range(size // replication): recv_rank = rank_col + j * replication start_time(start_inner_timer) + + # start_time(start_inner_timer) nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") - start_time(start_inner_timer) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") + + # start_time(start_inner_timer) rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ nnz_col_ids[j].size(0), matb._indices().size(1)) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskindalc-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskindalc-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) matb_send_indices = matb._indices()[:, nnz_row_mask] - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskind-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskind-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) matb_send_indices = matb_send_indices.float() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-indfloat-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-indfloat-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskval-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskval-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") - matb_send_list[j] = matb_send + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") + start_time(start_inner_timer) if recv_rank != q: selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) - send_objs[j] = dist.isend(selected_rows_count, dst=recv_rank) + dist.isend(selected_rows_count, tag=0, dst=recv_rank) + dist.isend(matb_send, tag=1, dst=recv_rank) else: matb_select_recv = matb_send.clone() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") - start_time(start_inner_timer) - for obj in send_objs: - if obj is not None: - obj.wait() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-wait-{name}") else: start_time(start_inner_timer) selected_rows_count_recv = torch.cuda.IntTensor(1) stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") - dist.irecv(selected_rows_count_recv, src=q).wait() + + start_time(start_inner_timer) + dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowcounts-{name}") start_time(start_inner_timer) matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") - torch.cuda.nvtx.range_pop() + start_time(start_inner_timer) + dist.irecv(matb_select_recv, tag=1, src=q).wait() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") - # Send selected rows from block row of matb from nnz_col_ids - start_time(start_timer) - if rank == q: - send_objs = [None] * (size // replication) - for j in range(size // replication): - recv_rank = rank_col + j * replication - if recv_rank != q: - send_objs[j] = dist.isend(matb_send_list[j], dst=recv_rank) - if name == "prob": - saspgemm_data += matb_send.numel() - for obj in send_objs: - if obj is not None: - obj.wait() - else: - dist.irecv(matb_select_recv, src=q).wait() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzrows-{name}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") torch.cuda.nvtx.range_pop() matb_recv_indices = matb_select_recv[:2, :].long() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 8e6e8b5..977bfa0 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -692,7 +692,7 @@ void rowselect_coo_gpu(const at::Tensor& nnz_cols, const at::Tensor& rows, const int BLOCK_SIZE = 256; - int BLOCK_COUNT = std::ceil(nnz_col_count / ((float) BLOCK_SIZE)); + int BLOCK_COUNT = std::ceil(row_count / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); if (nnz_col_count == 0) { From 49627ba00cbbeaea5affbe76b4f49e6e98915939 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 25 Aug 2022 17:55:53 -0700 Subject: [PATCH 091/178] Use SA-SpGEMM for row selection --- cagnet/samplers/ladies.py | 35 ++++++++++++++++++++++++++++++----- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 12d3923..4320059 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -242,6 +242,20 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-send-nnzrows") start_time(start_timer) if rank == q: + + # Add timings to timing_dict for src process for timing printouts + if f"spgemm-recv-alloccount-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-alloccount-{name}"] = [] + + if f"spgemm-recv-rowcounts-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-rowcounts-{name}"] = [] + + if f"spgemm-recv-alloc-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-alloc-{name}"] = [] + + if f"spgemm-recv-rowdata-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-rowdata-{name}"] = [] + for j in range(size // replication): recv_rank = rank_col + j * replication start_time(start_inner_timer) @@ -287,6 +301,12 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") else: + if f"spgemm-send-ovhd-{name}" not in timing_dict: + timing_dict[f"spgemm-send-ovhd-{name}"] = [] + + if f"spgemm-send-isend-{name}" not in timing_dict: + timing_dict[f"spgemm-send-isend-{name}"] = [] + start_time(start_inner_timer) selected_rows_count_recv = torch.cuda.IntTensor(1) stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") @@ -424,8 +444,6 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - # p_num_indices, p_num_values = dist_spgemm15D(batches, adj_matrix, replication, rank, size, \ - # row_groups, col_groups, "prob", timing_dict) p_num_indices, p_num_values = dist_saspgemm15D(batches, adj_matrix, replication, rank, size, \ row_groups, col_groups, "prob", sa_masks, sa_recv_buff, timing_dict) @@ -671,8 +689,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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)) - sampled_indices, sampled_values = dist_spgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ - row_groups, col_groups, "rowsel", timing_dict) + sa_masks.fill_(0) + sa_recv_buff.fill_(0) + sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ + row_groups, col_groups, "rowsel", sa_masks, + sa_recv_buff, timing_dict) sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, size=(nnz * mb_count, node_count_total)) @@ -733,7 +754,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"{k} min: {min_time} max: {max_time} avg: {avg_time} med: {med_time}") dist.barrier() for k, v in timing_dict.items(): - print(f"{k} total_time: {sum(v)} avg_time {sum(v) / len(v)} len: {len(v)}") + 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)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") print(f"spgemm_data: {spgemm_data}") From 32f41b360ae08248778b3e2ad984430e4670af02 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 25 Aug 2022 19:14:59 -0700 Subject: [PATCH 092/178] Mask out 0's for selection scatter_add's --- cagnet/samplers/ladies.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 4320059..05b4bc4 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -510,7 +510,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() - timing_dict["filter_darts"].append(stop_time(start_timer, stop_timer)) + timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") @@ -566,11 +566,11 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-invsum") - # dart_hits_inv_sum = torch.sparse.sum(dart_hits_inv_mtx, dim=1)._values() dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) - scatterd_add_gpu(dart_hits_inv_sum, next_frontier._indices()[0, :], dart_hits_inv, \ - dart_hits_inv.size(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) @@ -605,11 +605,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-samplecount") - # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) # sampled_count.scatter_add_(0, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values()) - scatteri_add_gpu(sampled_count, next_frontier_tmp._indices()[0,:], next_frontier_tmp._values(), \ - next_frontier_tmp._nnz()) + 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) From be0a6d26c6d61ccff9ad38558f9e331f54952a31 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 5 Sep 2022 15:46:56 -0700 Subject: [PATCH 093/178] Uniquely selects columns for accuracy testing --- cagnet/samplers/ladies.py | 55 ++++++++++++++++++++++++++++++--------- tests/test_ladies.py | 4 ++- 2 files changed, 45 insertions(+), 14 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 05b4bc4..c06a18d 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -30,16 +30,21 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name): else: timing_dict[range_name].append(stop_time(start_timer, stop_timer)) -saspgemm_data = 0 -spgemm_data = 0 +saspgemm_prob_data = 0 +spgemm_prob_data = 0 +saspgemm_rowsel_data = 0 +spgemm_rowsel_data = 0 def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ name, timing_dict=None): - global spgemm_data + global spgemm_prob_data + global spgemm_rowsel_data if name == "prob": - spgemm_data = 0 + spgemm_prob_data = 0 + if name == "rowsel": + spgemm_rowsel_data = 0 start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -77,7 +82,9 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-bcast-data-{name}") if name == "prob": - spgemm_data += matb_recv_indices.numel() + matb_recv_values.numel() + spgemm_prob_data += matb_recv_indices.numel() + matb_recv_values.numel() + if name == "rowsel": + spgemm_rowsel_data += matb_recv_indices.numel() + matb_recv_values.numel() start_time(start_timer) am_partid = rank_col * (size // replication ** 2) + i @@ -163,8 +170,13 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ name, nnz_row_masks, matb_recv_buff, timing_dict=None): - global saspgemm_data - saspgemm_data = 0 + global saspgemm_prob_data + global saspgemm_rowsel_data + + if name == "prob": + saspgemm_prob_data = 0 + if name == "rowsel": + saspgemm_rowsel_data = 0 start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -295,6 +307,11 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups if recv_rank != q: selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) dist.isend(selected_rows_count, tag=0, dst=recv_rank) + + if name == "prob": + saspgemm_prob_data += matb_send.numel() + if name == "rowsel": + saspgemm_rowsel_data += matb_send.numel() dist.isend(matb_send, tag=1, dst=recv_rank) else: matb_select_recv = matb_send.clone() @@ -525,11 +542,13 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-count-samples") # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() sampled_count = sampled_count.fill_(0) + 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) - scatteri_add_gpu(sampled_count, next_frontier_tmp._indices()[0, :], next_frontier_values, \ - next_frontier_values.size(0)) + sampled_count.scatter_add_(0, next_frontier_nnzidxs, next_frontier_nnzvals) torch.cuda.nvtx.range_pop() - timing_dict["count_samples"].append(stop_time(start_timer, stop_timer)) + timing_dict["count-samples"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-preproc") @@ -681,7 +700,15 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota col_select_mtx_rows = repeated_next_frontier.view(-1) col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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_(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)) @@ -762,6 +789,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"{k} total_time: {sum(v)} avg_time {avg_time} len: {len(v)}") print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") - print(f"spgemm_data: {spgemm_data}") - print(f"saspgemm_data: {saspgemm_data}") + print(f"spgemm_prob_data: {spgemm_prob_data}") + print(f"spgemm_rowsel_data: {spgemm_rowsel_data}") + print(f"saspgemm_prob_data: {saspgemm_prob_data}") + print(f"saspgemm_rowsel_data: {saspgemm_rowsel_data}") return batches_select, next_frontier_select, adj_matrices diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 828a4cc..795f931 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -82,7 +82,9 @@ def main(args): 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 + # 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 = [] From 03dae44a5ac3600c8e0d9a831d9153af418f77c2 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 10 Sep 2022 16:19:02 -0700 Subject: [PATCH 094/178] Verfied Cora accuracy p=1,2,4 --- examples/gcn_15d.py | 45 +++++++++++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 69a40a1..48c96c9 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -245,7 +245,7 @@ def evaluate(model, graph, features, labels, nid, nid_count, ampbyp, group): return correct.item() * 1.0 / nid_count # return correct.item() * 1.0 / len(labels) -def main(args): +def main(args, batches=None): # load and preprocess dataset device = torch.device('cuda') @@ -295,7 +295,8 @@ def main(args): os.environ["MASTER_PORT"] = "1234" print(f"hostname: {socket.gethostname()}", flush=True) - dist.init_process_group(backend=args.dist_backend) + 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) @@ -339,17 +340,20 @@ def main(args): train_nid = data.train_mask.nonzero().squeeze() test_nid = data.test_mask.nonzero().squeeze() - batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major torch.cuda.nvtx.range_pop() - torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - torch.manual_seed(0) - 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() + if batches is None: + print("beginning constructing batches") + batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major + torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") + # torch.manual_seed(0) + 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() + print("done constructing batches") batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) @@ -360,8 +364,9 @@ def main(args): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) batches_values = torch.cuda.DoubleTensor(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) + # g_loc = torch.pow(g_loc, 2) torch.cuda.nvtx.range_pop() + print(f"g_loc: {g_loc}") node_count = inputs.size(0) adj_matrix = adj_matrix.cuda() @@ -374,23 +379,23 @@ def main(args): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - torch.manual_seed(0) + # torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) - current_frontier, next_frontier, adj_matrices = \ - ladies_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + # current_frontier, next_frontier, adj_matrices = \ + # ladies_sampler(g_loc, batches_loc, args.batch_size, \ + # args.samp_num, args.n_bulkmb, \ + # args.n_layers, args.n_darts, \ + # args.replication, nnz_row_masks, sa_recv_buff, \ + # rank, size, row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - torch.manual_seed(0) + # torch.manual_seed(0) current_frontier, next_frontier, adj_matrices_bulk = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ From 75a97f1e4fdd42a6b31a066130265ce5c79f257f Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 12 Sep 2022 13:33:20 -0700 Subject: [PATCH 095/178] 1.5D dart throwing --- cagnet/samplers/ladies.py | 212 ++++++++++++++++---------- examples/gcn_15d.py | 22 +-- sparse-extension/sparse_coo_tensor.cu | 8 + 3 files changed, 148 insertions(+), 94 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index c06a18d..1acc4fd 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -445,6 +445,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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 @@ -511,13 +519,15 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") - dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() + # dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() + dart_values = torch.cuda.DoubleTensor(n_darts_col * mb_count).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") - compute_darts1d_gpu(dart_values, n_darts, mb_count) + # compute_darts1d_gpu(dart_values, n_darts, mb_count) + compute_darts1d_gpu(dart_values, n_darts_col, mb_count) torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) @@ -525,12 +535,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ - n_darts * mb_count, p._nnz()) + n_darts_col * mb_count, p._nnz()) + # n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") + 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, @@ -559,91 +571,116 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(select_start_timer) torch.cuda.nvtx.range_push("nvtx-dart-selection") - while overflowed_minibatches: - start_time(select_iter_start_timer) - torch.cuda.nvtx.range_push("nvtx-selection-iter") - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-psoverflow") - selection_iter_count += 1 - ps_overflow = torch.cumsum(overflow, dim=0) - total_overflow = ps_overflow[-1].item() - timing_dict["select-psoverflow"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-reciprocal") - 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) - torch.cuda.nvtx.range_push("nvtx-select-instmtx") - dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), - values=dart_hits_inv, - size=(mb_count, node_count_total)) - timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-invsum") - dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) - # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) - 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) - torch.cuda.nvtx.range_push("nvtx-select-psinv") - 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) - torch.cuda.nvtx.range_push("nvtx-select-computedarts") - dart_select = torch.cuda.DoubleTensor(total_overflow).uniform_() - # Compute darts for selection - compute_darts_select_gpu(dart_select, dart_hits_inv_sum, ps_dart_hits_inv_sum, ps_overflow, - mb_count, total_overflow) - timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-throwdarts") - # 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, - next_frontier._nnz()) - timing_dict["select-throwdarts"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") - 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=(mb_count, node_count_total)) - timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-samplecount") - sampled_count = torch.cuda.IntTensor(mb_count).fill_(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) - torch.cuda.nvtx.range_push("nvtx-select-overflow") - 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)) + if rank_col == 0: + while overflowed_minibatches: + start_time(select_iter_start_timer) + torch.cuda.nvtx.range_push("nvtx-selection-iter") + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-psoverflow") + 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) + torch.cuda.nvtx.range_push("nvtx-select-reciprocal") + 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) + torch.cuda.nvtx.range_push("nvtx-select-instmtx") + dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=dart_hits_inv, + size=(mb_count, node_count_total)) + timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-invsum") + dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) + # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) + 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) + torch.cuda.nvtx.range_push("nvtx-select-psinv") + 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) + torch.cuda.nvtx.range_push("nvtx-select-computedarts") + # 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, + mb_count, n_darts_select) + # mb_count, total_overflow) + timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-throwdarts") + # 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) + torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") + # 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=(mb_count, node_count_total)) + timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-samplecount") + sampled_count = torch.cuda.IntTensor(mb_count).fill_(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) + torch.cuda.nvtx.range_push("nvtx-select-overflow") + 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]) torch.cuda.nvtx.range_pop() # nvtx-dart-selection timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) - next_frontier_values = torch.logical_or(dart_hits_count, next_frontier._values()).long() + # 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=(mb_count, node_count_total)) @@ -661,10 +698,17 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-bool") + next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() + next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] + sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ + next_frontier._values()[next_frontier_nnzvals].int()) underfull_minibatches = (sampled_count < frontier_size).any() torch.cuda.nvtx.range_pop() timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + torch.cuda.nvtx.range_pop() # nvtx-sampling-iter timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 48c96c9..ac556b9 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -294,6 +294,7 @@ def main(args, batches=None): 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) @@ -346,7 +347,7 @@ def main(args, batches=None): print("beginning constructing batches") batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - # torch.manual_seed(0) + torch.manual_seed(0) vertex_perm = torch.randperm(train_nid.size(0)) # Generate minibatch vertices for i in range(args.n_bulkmb): @@ -364,7 +365,7 @@ def main(args, batches=None): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) batches_values = torch.cuda.DoubleTensor(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) + g_loc = torch.pow(g_loc, 2) torch.cuda.nvtx.range_pop() print(f"g_loc: {g_loc}") @@ -379,23 +380,24 @@ def main(args, batches=None): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - # torch.manual_seed(0) + torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) - # current_frontier, next_frontier, adj_matrices = \ - # ladies_sampler(g_loc, batches_loc, args.batch_size, \ - # args.samp_num, args.n_bulkmb, \ - # args.n_layers, args.n_darts, \ - # args.replication, nnz_row_masks, sa_recv_buff, \ - # rank, size, row_groups, col_groups) + torch.manual_seed(rank_col) + current_frontier, next_frontier, adj_matrices = \ + ladies_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - # torch.manual_seed(0) + torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices_bulk = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 977bfa0..bb6a301 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -442,6 +442,10 @@ void compute_darts_select_gpu(const at::Tensor& dart_select, 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(), @@ -544,6 +548,10 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, 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(), From 05b8c644a227f72b95b714534cd91fcea549f921 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 23 Sep 2022 13:53:01 -0700 Subject: [PATCH 096/178] Readded squaring probability values --- cagnet/samplers/ladies.py | 17 +++++++++++------ examples/gcn_15d.py | 25 +++++++++++++------------ sparse-extension/sparse_coo_tensor.cu | 12 +++++++++--- 3 files changed, 33 insertions(+), 21 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 1acc4fd..e5b5b99 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -247,14 +247,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups else: nnz_col_ids[j] = nnz_cols else: - dist.isend(nnz_cols, dst=q).wait() + dist.isend(nnz_cols, dst=q) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") # Send selected rows count torch.cuda.nvtx.range_push("nvtx-send-nnzrows") start_time(start_timer) if rank == q: - # Add timings to timing_dict for src process for timing printouts if f"spgemm-recv-alloccount-{name}" not in timing_dict: timing_dict[f"spgemm-recv-alloccount-{name}"] = [] @@ -316,7 +315,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups else: matb_select_recv = matb_send.clone() stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") - else: if f"spgemm-send-ovhd-{name}" not in timing_dict: timing_dict[f"spgemm-send-ovhd-{name}"] = [] @@ -479,13 +477,14 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) - # p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) + p_num_values = torch.square(p_num_values) scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + # p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) - print(f"p.nnz: {p._nnz()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) + print(f"p.nnz: {p._nnz()}") timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) @@ -505,6 +504,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True + p_rowsum = torch.cuda.DoubleTensor(mb_count).fill_(0) while underfull_minibatches: iter_count += 1 start_time(sample_start_timer) @@ -514,6 +514,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota torch.cuda.nvtx.range_push("nvtx-prob-rowsum") ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) ps_p_values[0] = 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 torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) @@ -527,7 +530,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-dart-throw") # compute_darts1d_gpu(dart_values, n_darts, mb_count) - compute_darts1d_gpu(dart_values, n_darts_col, mb_count) + # compute_darts1d_gpu(dart_values, n_darts_col, mb_count) + compute_darts1d_gpu(dart_values, p_rowsum, ps_p_rowsum, n_darts, mb_count) + torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index ac556b9..41c8b22 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -329,6 +329,7 @@ def main(args, batches=None): print("normalizing", flush=True) 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)) @@ -347,7 +348,7 @@ def main(args, batches=None): print("beginning constructing batches") batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - torch.manual_seed(0) + # torch.manual_seed(0) vertex_perm = torch.randperm(train_nid.size(0)) # Generate minibatch vertices for i in range(args.n_bulkmb): @@ -365,7 +366,7 @@ def main(args, batches=None): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) batches_values = torch.cuda.DoubleTensor(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) + # g_loc = torch.pow(g_loc, 2) torch.cuda.nvtx.range_pop() print(f"g_loc: {g_loc}") @@ -380,19 +381,19 @@ def main(args, batches=None): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - torch.manual_seed(0) + # torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) - torch.manual_seed(rank_col) - current_frontier, next_frontier, adj_matrices = \ - ladies_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + # torch.manual_seed(rank_col) + # current_frontier, next_frontier, adj_matrices = \ + # ladies_sampler(g_loc, batches_loc, args.batch_size, \ + # args.samp_num, args.n_bulkmb, \ + # args.n_layers, args.n_darts, \ + # args.replication, nnz_row_masks, sa_recv_buff, \ + # rank, size, row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() @@ -435,8 +436,8 @@ def main(args, batches=None): torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0), size=(node_count, node_count)) adj_matrix = adj_matrix.coalesce().cuda() - adj_matrix = row_normalize(adj_matrix) - adj_matrix = torch.pow(adj_matrix, 2) + # adj_matrix = row_normalize(adj_matrix) + # adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code return current_frontier, next_frontier, adj_matrices, adj_matrix diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index bb6a301..f45a2df 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -393,7 +393,8 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } -__global__ void ComputeDarts1D(double *dart_values, int n_darts, int mb_count) { +__global__ void ComputeDarts1D(double *dart_values, double *p_rowsum, double *ps_p_rowsum, + int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -401,17 +402,22 @@ __global__ void ComputeDarts1D(double *dart_values, int n_darts, int mb_count) { int dart_count = n_darts * mb_count; for (int i = id; i < dart_count; i += stride) { int row = i / n_darts; - dart_values[i] += (double)row; + // dart_values[i] += (double)row; + dart_values[i] *= p_rowsum[row]; + dart_values[i] += ps_p_rowsum[row]; } } -void compute_darts1d_gpu(const at::Tensor& dart_values, int n_darts, int mb_count) { +void compute_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& p_rowsum, + const at::Tensor& ps_p_rowsum, 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); ComputeDarts1D<<>>(dart_values.data(), + p_rowsum.data(), + ps_p_rowsum.data(), n_darts, mb_count); CHECK_ERROR("dart1d computation error") From 5691bd32a0f02d78683c54a7ee0ded44902486f7 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 23 Sep 2022 19:21:57 -0700 Subject: [PATCH 097/178] Extracted API --- cagnet/samplers/ladies.py | 768 +------------------------------------ cagnet/samplers/utils.py | 770 ++++++++++++++++++++++++++++++++++++++ examples/gcn_15d.py | 18 +- 3 files changed, 788 insertions(+), 768 deletions(-) create mode 100644 cagnet/samplers/utils.py diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index e5b5b99..9a6d352 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -3,6 +3,7 @@ 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, \ @@ -18,420 +19,9 @@ def stop_time(start_timer, stop_timer): torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) -def stop_time_add(start_timer, stop_timer, timing_dict, range_name): - if timing_dict is not None: - # misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ - # not range_name.startswith("spgemm-bcast-data") and \ - # not range_name.startswith("spgemm-local-spgemm") and \ - # not range_name.startswith("spgemm-reduce-nnz") and \ - # not range_name.startswith("spgemm-reduce") - if False and misc_range: - timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) - else: - timing_dict[range_name].append(stop_time(start_timer, stop_timer)) - -saspgemm_prob_data = 0 -spgemm_prob_data = 0 -saspgemm_rowsel_data = 0 -spgemm_rowsel_data = 0 - -def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ - name, timing_dict=None): - - global spgemm_prob_data - global spgemm_rowsel_data - - if name == "prob": - spgemm_prob_data = 0 - if name == "rowsel": - spgemm_rowsel_data = 0 - - 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}") - - 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}") - - 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}") - - 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}") - - if name == "prob": - spgemm_prob_data += matb_recv_indices.numel() + matb_recv_values.numel() - if name == "rowsel": - spgemm_rowsel_data += matb_recv_indices.numel() + matb_recv_values.numel() - - 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}") - - start_time(start_timer) - 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) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") - - start_time(start_timer) - matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") - - start_time(start_timer) - matc_chunk = matc_chunk.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") - - start_time(start_timer) - matc += matc_chunk - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") - - start_time(start_timer) - matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") - - # 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}") - - 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}") - - 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}") - - 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())) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}") - - 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}") - start_time(start_timer) - matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") - - 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}") - - return matc_nnz_indices, matc_nnz_values - -def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ - name, nnz_row_masks, matb_recv_buff, timing_dict=None): - - global saspgemm_prob_data - global saspgemm_rowsel_data - - if name == "prob": - saspgemm_prob_data = 0 - if name == "rowsel": - saspgemm_rowsel_data = 0 - - 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))) - 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}") - - 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 - 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}") - - # Determine number of nonzero columns in chunk of mata - start_time(start_timer) - nnz_cols = torch.unique(mata_chunk_indices[1,:]) - 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}") - - # Gather nnz column counts on rank q - start_time(start_timer) - if rank == q: - nnz_cols_count_list = [] - for j in range(size // replication): - nnz_cols_count_list.append(torch.cuda.IntTensor(1).fill_(0)) - 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}") - - # 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}") - - # 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.irecv(nnz_col_ids[j], src=recv_rank)) - else: - nnz_col_ids[j] = nnz_cols - else: - dist.isend(nnz_cols, dst=q) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") - - # Send selected rows count - torch.cuda.nvtx.range_push("nvtx-send-nnzrows") - start_time(start_timer) - if rank == q: - # Add timings to timing_dict for src process for timing printouts - if f"spgemm-recv-alloccount-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-alloccount-{name}"] = [] - - if f"spgemm-recv-rowcounts-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-rowcounts-{name}"] = [] - - if f"spgemm-recv-alloc-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-alloc-{name}"] = [] - - if f"spgemm-recv-rowdata-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-rowdata-{name}"] = [] - - for j in range(size // replication): - recv_rank = rank_col + j * replication - start_time(start_inner_timer) - - # start_time(start_inner_timer) - nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") - - # start_time(start_inner_timer) - rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ - nnz_col_ids[j].size(0), matb._indices().size(1)) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") - - # start_time(start_inner_timer) - matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskindalc-{name}") - - # start_time(start_inner_timer) - matb_send_indices = matb._indices()[:, nnz_row_mask] - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskind-{name}") - - # start_time(start_inner_timer) - matb_send_indices = matb_send_indices.float() - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-indfloat-{name}") - - # start_time(start_inner_timer) - matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskval-{name}") - - # start_time(start_inner_timer) - matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") - - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") - - start_time(start_inner_timer) - if recv_rank != q: - selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) - dist.isend(selected_rows_count, tag=0, dst=recv_rank) - - if name == "prob": - saspgemm_prob_data += matb_send.numel() - if name == "rowsel": - saspgemm_rowsel_data += matb_send.numel() - dist.isend(matb_send, tag=1, dst=recv_rank) - else: - matb_select_recv = matb_send.clone() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") - else: - if f"spgemm-send-ovhd-{name}" not in timing_dict: - timing_dict[f"spgemm-send-ovhd-{name}"] = [] - - if f"spgemm-send-isend-{name}" not in timing_dict: - timing_dict[f"spgemm-send-isend-{name}"] = [] - - start_time(start_inner_timer) - selected_rows_count_recv = torch.cuda.IntTensor(1) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") - - start_time(start_inner_timer) - dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowcounts-{name}") - - start_time(start_inner_timer) - matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") - - start_time(start_inner_timer) - dist.irecv(matb_select_recv, tag=1, src=q).wait() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") - - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") - torch.cuda.nvtx.range_pop() - - matb_recv_indices = matb_select_recv[:2, :].long() - matb_recv_values = matb_select_recv[2, :].double() - - start_time(start_timer) - 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) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") - - start_time(start_timer) - matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") - - start_time(start_timer) - matc_chunk = matc_chunk.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") - - start_time(start_timer) - matc += matc_chunk - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") - - start_time(start_timer) - matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") - - # 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}") - - 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}") - - 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}") - - 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())) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}") - - 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}") - start_time(start_timer) - matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") - - 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}") - - return matc_nnz_indices, matc_nnz_values - def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups): - global spgemm_data - global saspgemm_data - - 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) - total_start_timer = torch.cuda.Event(enable_timing=True) total_stop_timer = torch.cuda.Event(enable_timing=True) @@ -461,356 +51,22 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota else: nnz = current_frontier[0, :].size(0) - # A * Q^i - # indices would change based on mb_count - - # TODO: assume n_layers=1 for now - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = dist_saspgemm15D(batches, adj_matrix, replication, rank, size, \ - row_groups, col_groups, "prob", sa_masks, sa_recv_buff, - timing_dict) - - torch.cuda.nvtx.range_pop() - timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-compute-p") - p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) - p_num_values = torch.square(p_num_values) - scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) - # p_den = p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) - p = torch.sparse_coo_tensor(indices=p_num_indices, - values=p_num_values, - size=(mb_count, node_count_total)) - normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) - print(f"p.nnz: {p._nnz()}") - timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-pre-loop") - next_frontier = torch.sparse_coo_tensor(indices=p._indices(), - values=torch.cuda.LongTensor(p._nnz()).fill_(0), - size=(mb_count, node_count_total)) - sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) - - torch.cuda.nvtx.range_pop() - timing_dict["pre-loop"].append(stop_time(start_timer, stop_timer)) - - iter_count = 0 - selection_iter_count = 0 - torch.cuda.nvtx.range_push("nvtx-sampling") - - # underfull_minibatches = (sampled_count < frontier_size).any() - underfull_minibatches = True - - p_rowsum = torch.cuda.DoubleTensor(mb_count).fill_(0) - while underfull_minibatches: - iter_count += 1 - start_time(sample_start_timer) - torch.cuda.nvtx.range_push("nvtx-sampling-iter") - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-prob-rowsum") - ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) - ps_p_values[0] = 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 - torch.cuda.nvtx.range_pop() - timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-gen-darts") - # dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() - dart_values = torch.cuda.DoubleTensor(n_darts_col * mb_count).uniform_() - torch.cuda.nvtx.range_pop() - timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-dart-throw") - # 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, mb_count) - - torch.cuda.nvtx.range_pop() - timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-filter-darts") - dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) - throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ - n_darts_col * mb_count, p._nnz()) - # n_darts * mb_count, p._nnz()) - torch.cuda.nvtx.range_pop() - timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-add-to-frontier") - 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=(mb_count, node_count_total)) - torch.cuda.nvtx.range_pop() - timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-count-samples") - # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() - sampled_count = sampled_count.fill_(0) - 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) - torch.cuda.nvtx.range_pop() - timing_dict["count-samples"].append(stop_time(start_timer, stop_timer)) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-preproc") - overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - overflowed_minibatches = (overflow > 0).any() - timing_dict["select-preproc"].append(stop_time(start_timer, stop_timer)) + p = gen_prob_dist(batches, adj_matrix, mb_count, node_count_total, replication, rank, size, \ + row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict) - start_time(select_start_timer) - torch.cuda.nvtx.range_push("nvtx-dart-selection") - if rank_col == 0: - while overflowed_minibatches: - start_time(select_iter_start_timer) - torch.cuda.nvtx.range_push("nvtx-selection-iter") + next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ + row_groups, col_groups, timing_dict) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-psoverflow") - 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)) + batches_select, next_frontier_select, adj_matrix_sample = \ + select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ + batch_size, frontier_size, mb_count, mb_count_total, node_count_total, \ + replication, rank, size, row_groups, col_groups, timing_dict, i) - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-reciprocal") - 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) - torch.cuda.nvtx.range_push("nvtx-select-instmtx") - dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), - values=dart_hits_inv, - size=(mb_count, node_count_total)) - timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-invsum") - dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) - # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) - 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) - torch.cuda.nvtx.range_push("nvtx-select-psinv") - 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) - torch.cuda.nvtx.range_push("nvtx-select-computedarts") - # 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, - mb_count, n_darts_select) - # mb_count, total_overflow) - timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-throwdarts") - # 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) - torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") - # 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=(mb_count, node_count_total)) - timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-samplecount") - sampled_count = torch.cuda.IntTensor(mb_count).fill_(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) - torch.cuda.nvtx.range_push("nvtx-select-overflow") - 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]) - - torch.cuda.nvtx.range_pop() # nvtx-dart-selection - timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) - - # 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=(mb_count, node_count_total)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-probs") - 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,:]) - p._values()[filled_minibatches_mask] = 0 - torch.cuda.nvtx.range_pop() - timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-compute-bool") - next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() - next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] - sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) - sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ - next_frontier._values()[next_frontier_nnzvals].int()) - underfull_minibatches = (sampled_count < frontier_size).any() - torch.cuda.nvtx.range_pop() - timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) - - overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - - torch.cuda.nvtx.range_pop() # nvtx-sampling-iter - timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) - - torch.cuda.nvtx.range_pop() # nvtx-sampling - - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-construct-nextf") - 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) - next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) - torch.cuda.nvtx.range_pop() - timing_dict["construct-nextf"].append(stop_time(start_timer, stop_timer)) - - 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 i == 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() - scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() - repeated_next_frontier.add_(scale_mtx[:, None]) - col_select_mtx_rows = repeated_next_frontier.view(-1) - col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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)) - - # 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)) - sa_masks.fill_(0) - sa_recv_buff.fill_(0) - sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ - row_groups, col_groups, "rowsel", sa_masks, - sa_recv_buff, timing_dict) - sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, - size=(nnz * mb_count, node_count_total)) - - torch.cuda.nvtx.range_pop() - timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) - - # 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) - shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - rank, size, replication, batch_size, node_count_total, mb_count_total) - sampled_indices[1,:] += row_shift - torch.cuda.nvtx.range_pop() - timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) - - # 4. Multiply sampled rows with col_select matrix - start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, - size=(nnz * mb_count, node_count_total * mb_count_total)) - 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))) - sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ - row_groups, col_groups, "colsel", timing_dict) - torch.cuda.nvtx.range_pop() - timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) - - # 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_matrices[i] = adj_matrix_sample - current_frontier = next_frontier - torch.cuda.nvtx.range_pop() - timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) - torch.cuda.nvtx.range_pop() print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) for k, v in sorted(timing_dict.items()): @@ -836,10 +92,4 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota else: avg_time = -1.0 print(f"{k} total_time: {sum(v)} avg_time {avg_time} len: {len(v)}") - print(f"iter_count: {iter_count}") - print(f"selection_iter_count: {selection_iter_count}") - print(f"spgemm_prob_data: {spgemm_prob_data}") - print(f"spgemm_rowsel_data: {spgemm_rowsel_data}") - print(f"saspgemm_prob_data: {saspgemm_prob_data}") - print(f"saspgemm_rowsel_data: {saspgemm_rowsel_data}") return batches_select, next_frontier_select, adj_matrices diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py new file mode 100644 index 0000000..b362529 --- /dev/null +++ b/cagnet/samplers/utils.py @@ -0,0 +1,770 @@ +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, \ + shift_rowselect_gpu, shift_colselect_gpu, \ + scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu + +def start_time(timer): + timer.record() + +def stop_time(start_timer, stop_timer): + stop_timer.record() + torch.cuda.synchronize() + return start_timer.elapsed_time(stop_timer) + +def stop_time_add(start_timer, stop_timer, timing_dict, range_name): + if timing_dict is not None: + # misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ + # not range_name.startswith("spgemm-bcast-data") and \ + # not range_name.startswith("spgemm-local-spgemm") and \ + # not range_name.startswith("spgemm-reduce-nnz") and \ + # not range_name.startswith("spgemm-reduce") + if False and misc_range: + timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) + else: + timing_dict[range_name].append(stop_time(start_timer, stop_timer)) + + +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}") + + 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}") + + 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}") + + 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}") + + 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}") + + start_time(start_timer) + 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) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + + start_time(start_timer) + matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + + start_time(start_timer) + matc_chunk = matc_chunk.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + + start_time(start_timer) + matc += matc_chunk + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") + + # 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}") + + 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}") + + 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}") + + 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())) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}") + + 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}") + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") + + 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}") + + return matc_nnz_indices, matc_nnz_values + +def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ + name, nnz_row_masks, matb_recv_buff, timing_dict=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))) + 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}") + + 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 + 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}") + + # Determine number of nonzero columns in chunk of mata + start_time(start_timer) + nnz_cols = torch.unique(mata_chunk_indices[1,:]) + 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}") + + # Gather nnz column counts on rank q + start_time(start_timer) + if rank == q: + nnz_cols_count_list = [] + for j in range(size // replication): + nnz_cols_count_list.append(torch.cuda.IntTensor(1).fill_(0)) + 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}") + + # 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}") + + # 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.irecv(nnz_col_ids[j], src=recv_rank)) + else: + nnz_col_ids[j] = nnz_cols + else: + dist.isend(nnz_cols, dst=q) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") + + # Send selected rows count + torch.cuda.nvtx.range_push("nvtx-send-nnzrows") + start_time(start_timer) + if rank == q: + # Add timings to timing_dict for src process for timing printouts + if f"spgemm-recv-alloccount-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-alloccount-{name}"] = [] + + if f"spgemm-recv-rowcounts-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-rowcounts-{name}"] = [] + + if f"spgemm-recv-alloc-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-alloc-{name}"] = [] + + if f"spgemm-recv-rowdata-{name}" not in timing_dict: + timing_dict[f"spgemm-recv-rowdata-{name}"] = [] + + for j in range(size // replication): + recv_rank = rank_col + j * replication + start_time(start_inner_timer) + + # start_time(start_inner_timer) + nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") + + # start_time(start_inner_timer) + rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ + nnz_col_ids[j].size(0), matb._indices().size(1)) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") + + # start_time(start_inner_timer) + matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskindalc-{name}") + + # start_time(start_inner_timer) + matb_send_indices = matb._indices()[:, nnz_row_mask] + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskind-{name}") + + # start_time(start_inner_timer) + matb_send_indices = matb_send_indices.float() + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-indfloat-{name}") + + # start_time(start_inner_timer) + matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskval-{name}") + + # start_time(start_inner_timer) + matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") + + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") + + start_time(start_inner_timer) + if recv_rank != q: + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) + dist.isend(selected_rows_count, tag=0, dst=recv_rank) + dist.isend(matb_send, tag=1, dst=recv_rank) + else: + matb_select_recv = matb_send.clone() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") + else: + if f"spgemm-send-ovhd-{name}" not in timing_dict: + timing_dict[f"spgemm-send-ovhd-{name}"] = [] + + if f"spgemm-send-isend-{name}" not in timing_dict: + timing_dict[f"spgemm-send-isend-{name}"] = [] + + start_time(start_inner_timer) + selected_rows_count_recv = torch.cuda.IntTensor(1) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") + + start_time(start_inner_timer) + dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowcounts-{name}") + + start_time(start_inner_timer) + matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") + + start_time(start_inner_timer) + dist.irecv(matb_select_recv, tag=1, src=q).wait() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") + + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") + torch.cuda.nvtx.range_pop() + + matb_recv_indices = matb_select_recv[:2, :].long() + matb_recv_values = matb_select_recv[2, :].double() + + start_time(start_timer) + 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) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + + start_time(start_timer) + matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + + start_time(start_timer) + matc_chunk = matc_chunk.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + + start_time(start_timer) + matc += matc_chunk + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") + + # 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}") + + 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}") + + 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}") + + 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())) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-reduce-{name}") + + 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}") + start_time(start_timer) + matc = matc.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coalesce-{name}") + + 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}") + + 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, sa_recv_buff, timing_dict): + + 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) + torch.cuda.nvtx.range_push("nvtx-probability-spgemm") + p_num_indices, p_num_values = dist_saspgemm15D(numerator, adj_matrix, replication, rank, size, \ + row_groups, col_groups, "prob", sa_masks, sa_recv_buff, + timing_dict) + + torch.cuda.nvtx.range_pop() + timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-compute-p") + p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) + p_num_values = torch.square(p_num_values) + scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) + print(f"p.nnz: {p._nnz()}") + normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) + 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): + + 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 + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-pre-loop") + next_frontier = torch.sparse_coo_tensor(indices=p._indices(), + values=torch.cuda.LongTensor(p._nnz()).fill_(0), + size=(mb_count, node_count_total)) + sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + + torch.cuda.nvtx.range_pop() + timing_dict["pre-loop"].append(stop_time(start_timer, stop_timer)) + + iter_count = 0 + selection_iter_count = 0 + torch.cuda.nvtx.range_push("nvtx-sampling") + + # underfull_minibatches = (sampled_count < frontier_size).any() + underfull_minibatches = True + + p_rowsum = torch.cuda.DoubleTensor(mb_count).fill_(0) + while underfull_minibatches: + iter_count += 1 + start_time(sample_start_timer) + torch.cuda.nvtx.range_push("nvtx-sampling-iter") + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-prob-rowsum") + ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) + ps_p_values[0] = 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 + torch.cuda.nvtx.range_pop() + timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-gen-darts") + # dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() + dart_values = torch.cuda.DoubleTensor(n_darts_col * mb_count).uniform_() + torch.cuda.nvtx.range_pop() + timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-dart-throw") + # 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, mb_count) + + torch.cuda.nvtx.range_pop() + timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-filter-darts") + dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) + throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ + n_darts_col * mb_count, p._nnz()) + # n_darts * mb_count, p._nnz()) + torch.cuda.nvtx.range_pop() + timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-add-to-frontier") + 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=(mb_count, node_count_total)) + torch.cuda.nvtx.range_pop() + timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-count-samples") + # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() + sampled_count = sampled_count.fill_(0) + 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) + torch.cuda.nvtx.range_pop() + timing_dict["count-samples"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-preproc") + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + overflowed_minibatches = (overflow > 0).any() + timing_dict["select-preproc"].append(stop_time(start_timer, stop_timer)) + + start_time(select_start_timer) + torch.cuda.nvtx.range_push("nvtx-dart-selection") + + if rank_col == 0: + while overflowed_minibatches: + start_time(select_iter_start_timer) + torch.cuda.nvtx.range_push("nvtx-selection-iter") + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-psoverflow") + 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) + torch.cuda.nvtx.range_push("nvtx-select-reciprocal") + 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) + torch.cuda.nvtx.range_push("nvtx-select-instmtx") + dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), + values=dart_hits_inv, + size=(mb_count, node_count_total)) + timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-invsum") + dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) + # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) + 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) + torch.cuda.nvtx.range_push("nvtx-select-psinv") + 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) + torch.cuda.nvtx.range_push("nvtx-select-computedarts") + # 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, + mb_count, n_darts_select) + # mb_count, total_overflow) + timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-throwdarts") + # 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) + torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") + # 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=(mb_count, node_count_total)) + timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-select-samplecount") + sampled_count = torch.cuda.IntTensor(mb_count).fill_(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) + torch.cuda.nvtx.range_push("nvtx-select-overflow") + 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]) + + torch.cuda.nvtx.range_pop() # nvtx-dart-selection + timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) + + # 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=(mb_count, node_count_total)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-set-probs") + 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,:]) + p._values()[filled_minibatches_mask] = 0 + torch.cuda.nvtx.range_pop() + timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-compute-bool") + next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() + next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] + sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ + next_frontier._values()[next_frontier_nnzvals].int()) + underfull_minibatches = (sampled_count < frontier_size).any() + torch.cuda.nvtx.range_pop() + timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) + + overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + + torch.cuda.nvtx.range_pop() # nvtx-sampling-iter + timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) + + torch.cuda.nvtx.range_pop() # nvtx-sampling + print(f"iter_count: {iter_count}") + print(f"selection_iter_count: {selection_iter_count}") + + return next_frontier + +def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ + batch_size, frontier_size, mb_count, mb_count_total, node_count_total, replication, \ + rank, size, row_groups, col_groups, timing_dict, layer_id): + + start_timer = torch.cuda.Event(enable_timing=True) + stop_timer = torch.cuda.Event(enable_timing=True) + + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-construct-nextf") + 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) + next_frontier_select = torch.cat((next_frontier_select, batches_select), dim=1) + torch.cuda.nvtx.range_pop() + timing_dict["construct-nextf"].append(stop_time(start_timer, stop_timer)) + + 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() + scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() + repeated_next_frontier.add_(scale_mtx[:, None]) + col_select_mtx_rows = repeated_next_frontier.view(-1) + col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).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)) + + # 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)) + sa_masks.fill_(0) + sa_recv_buff.fill_(0) + sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ + row_groups, col_groups, "rowsel", sa_masks, + sa_recv_buff, timing_dict) + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + size=(nnz * mb_count, node_count_total)) + + torch.cuda.nvtx.range_pop() + timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) + + # 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) + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, replication, batch_size, node_count_total, mb_count_total) + sampled_indices[1,:] += row_shift + torch.cuda.nvtx.range_pop() + timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) + + # 4. Multiply sampled rows with col_select matrix + start_time(start_timer) + torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + size=(nnz * mb_count, node_count_total * mb_count_total)) + 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))) + sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ + row_groups, col_groups, "colsel", timing_dict) + torch.cuda.nvtx.range_pop() + timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) + + # 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))) + + torch.cuda.nvtx.range_pop() + timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) + + torch.cuda.nvtx.range_pop() # nvtx-select-mtxs + + return batches_select, next_frontier_select, adj_matrix_sample diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 41c8b22..f09d0ef 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -348,7 +348,7 @@ def main(args, batches=None): print("beginning constructing batches") batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - # torch.manual_seed(0) + torch.manual_seed(0) vertex_perm = torch.randperm(train_nid.size(0)) # Generate minibatch vertices for i in range(args.n_bulkmb): @@ -381,19 +381,19 @@ def main(args, batches=None): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - # torch.manual_seed(0) + torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) - # torch.manual_seed(rank_col) - # current_frontier, next_frontier, adj_matrices = \ - # ladies_sampler(g_loc, batches_loc, args.batch_size, \ - # args.samp_num, args.n_bulkmb, \ - # args.n_layers, args.n_darts, \ - # args.replication, nnz_row_masks, sa_recv_buff, \ - # rank, size, row_groups, col_groups) + torch.manual_seed(rank_col) + current_frontier, next_frontier, adj_matrices = \ + ladies_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) print() torch.cuda.profiler.cudart().cudaProfilerStart() From 7617e07314a455c0c579114180a0e9eb793c12ce Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 27 Sep 2022 11:54:39 -0700 Subject: [PATCH 098/178] Preliminary SAGE impl --- cagnet/samplers/__init__.py | 2 + cagnet/samplers/ladies.py | 2 - cagnet/samplers/sage.py | 128 ++++++++++++++++++++++++++++++++++++ cagnet/samplers/utils.py | 92 ++++++++++++++++---------- examples/gcn_15d.py | 77 +++++++++++++++------- 5 files changed, 243 insertions(+), 58 deletions(-) create mode 100644 cagnet/samplers/sage.py diff --git a/cagnet/samplers/__init__.py b/cagnet/samplers/__init__.py index fd68782..79e0db9 100644 --- a/cagnet/samplers/__init__.py +++ b/cagnet/samplers/__init__.py @@ -1,7 +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 index 9a6d352..9d2e82f 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -51,11 +51,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota else: nnz = current_frontier[0, :].size(0) - p = gen_prob_dist(batches, adj_matrix, mb_count, node_count_total, replication, rank, size, \ row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict) - next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ row_groups, col_groups, timing_dict) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py new file mode 100644 index 0000000..3be91d3 --- /dev/null +++ b/cagnet/samplers/sage.py @@ -0,0 +1,128 @@ +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, \ + scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu + +def start_time(timer): + timer.record() + +def stop_time(start_timer, stop_timer): + stop_timer.record() + torch.cuda.synchronize() + return start_timer.elapsed_time(stop_timer) + +def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ + replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups): + + total_start_timer = torch.cuda.Event(enable_timing=True) + total_stop_timer = torch.cuda.Event(enable_timing=True) + + timing_dict = defaultdict(list) + + 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 mb mtx for layer j + + start_time(total_start_timer) + for i in range(n_layers): + if i == 0: + nnz = batch_size + else: + nnz = current_frontier[0, :].size(0) + + # Expand batches matrix + batches_expand_rows = torch.arange(mb_count * nnz, \ + device=torch.device("cuda:0")) + 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 * nnz, node_count_total)) + + p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, + replication, rank, size, row_groups, col_groups, + sa_masks, sa_recv_buff, timing_dict, "sage") + + next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, + replication, rank, size, row_groups, col_groups, + timing_dict) + + # Collapse next_frontier matrix + next_frontier_nnz = next_frontier._values().nonzero().squeeze() + frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], + bins=p.size(0)).int() + + frontier_nnz_sizes = torch.clamp(frontier_nnz_sizes, max=frontier_size) + print(f"frontier_nnz_sizes: {frontier_nnz_sizes}") + frontier_nnz_sizes = frontier_nnz_sizes.sum() # just for mb=1 + next_frontier_rows = torch.repeat_interleave( + torch.arange(mb_count, device=torch.device("cuda:0")), + frontier_nnz_sizes) + next_frontier_idxs = torch.stack( + (next_frontier_rows, + next_frontier._indices()[1, next_frontier_nnz])) + next_frontier = torch.sparse_coo_tensor(next_frontier_idxs, + next_frontier._values()[next_frontier_nnz], + size=(mb_count, node_count_total)) + next_frontier = next_frontier.coalesce() + next_frontier._values().fill_(1) + + batches_select, next_frontier_select, adj_matrix_sample = \ + select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, + nnz, batch_size, batch_size * frontier_size, mb_count, + mb_count_total, node_count_total, replication, rank, + size, row_groups, col_groups, timing_dict, i) + + adj_matrices[i] = adj_matrix_sample + current_frontier = next_frontier + + + print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) + 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": + # 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/utils.py b/cagnet/samplers/utils.py index b362529..24942a3 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -383,8 +383,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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, sa_recv_buff, timing_dict): +def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication, + rank, size, row_groups, col_groups, sa_masks, sa_recv_buff, + timing_dict, name): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -392,26 +393,32 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # TODO: assume n_layers=1 for now start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-probability-spgemm") - p_num_indices, p_num_values = dist_saspgemm15D(numerator, adj_matrix, replication, rank, size, \ - row_groups, col_groups, "prob", sa_masks, sa_recv_buff, - timing_dict) + p_num_indices, p_num_values = dist_saspgemm15D( + numerator, adj_matrix, replication, rank, size, + row_groups, col_groups, "prob", sa_masks, + sa_recv_buff, timing_dict) torch.cuda.nvtx.range_pop() timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p_den = torch.cuda.DoubleTensor(mb_count).fill_(0) - p_num_values = torch.square(p_num_values) + p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) + if name == "ladies": + p_num_values = torch.square(p_num_values) + elif name == "sage": + p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) - p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(mb_count, node_count_total)) + p = torch.sparse_coo_tensor(indices=p_num_indices, + values=p_num_values, + size=(numerator.size(0), node_count_total)) print(f"p.nnz: {p._nnz()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) 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): +def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, + rank, size, row_groups, col_groups, timing_dict): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -436,9 +443,13 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") next_frontier = torch.sparse_coo_tensor(indices=p._indices(), - values=torch.cuda.LongTensor(p._nnz()).fill_(0), - size=(mb_count, node_count_total)) - sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + values=torch.cuda.LongTensor(p._nnz()).fill_(0), + size=(p.size(0), node_count_total)) + + frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], + bins=p.size(0)).int() + # sampled_count = torch.cuda.IntTensor(p.size(0)).fill_(0) + sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) torch.cuda.nvtx.range_pop() timing_dict["pre-loop"].append(stop_time(start_timer, stop_timer)) @@ -450,26 +461,31 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True - p_rowsum = torch.cuda.DoubleTensor(mb_count).fill_(0) + p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) while underfull_minibatches: iter_count += 1 + print(f"iter_count: {iter_count}", flush=True) + print(f"sampled_count: {sampled_count}") start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-prob-rowsum") - ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) - ps_p_values[0] = 0 + # ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) + ps_p_values = torch.cumsum(p._values(), dim=0) + # 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 + # ps_p_rowsum = torch.cumsum(p_rowsum, dim=0).roll(1) + ps_p_rowsum = torch.cumsum(p_rowsum, dim=0) + # ps_p_rowsum[0] = 0 torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-gen-darts") # dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() - dart_values = torch.cuda.DoubleTensor(n_darts_col * mb_count).uniform_() + dart_values = torch.cuda.DoubleTensor(n_darts_col * p.size(0)).uniform_() torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) @@ -477,7 +493,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r torch.cuda.nvtx.range_push("nvtx-dart-throw") # 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, mb_count) + compute_darts1d_gpu(dart_values, p_rowsum, ps_p_rowsum, n_darts_col, p.size(0)) torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) @@ -486,7 +502,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r torch.cuda.nvtx.range_push("nvtx-filter-darts") dart_hits_count = torch.cuda.IntTensor(p._nnz()).fill_(0) throw_darts1d_gpu(dart_values, ps_p_values, dart_hits_count, \ - n_darts_col * mb_count, p._nnz()) + n_darts_col * p.size(0), p._nnz()) # n_darts * mb_count, p._nnz()) torch.cuda.nvtx.range_pop() timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) @@ -494,10 +510,12 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-add-to-frontier") 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_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=(mb_count, node_count_total)) + values=next_frontier_values, + size=(p.size(0), node_count_total)) torch.cuda.nvtx.range_pop() timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) @@ -505,6 +523,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r torch.cuda.nvtx.range_push("nvtx-count-samples") # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() sampled_count = sampled_count.fill_(0) + sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) 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] @@ -530,6 +549,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-psoverflow") selection_iter_count += 1 + print(f"selection_iter_count: {selection_iter_count}", flush=True) ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() # n_darts_select = total_overflow // replication @@ -548,12 +568,12 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r torch.cuda.nvtx.range_push("nvtx-select-instmtx") dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), values=dart_hits_inv, - size=(mb_count, node_count_total)) + size=(p.size(0), node_count_total)) timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-invsum") - dart_hits_inv_sum = torch.cuda.DoubleTensor(mb_count).fill_(0) + dart_hits_inv_sum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) dart_hits_inv_nnzidxs = dart_hits_inv.nonzero().squeeze() indices_nnz_idxs = next_frontier._indices()[0,dart_hits_inv_nnzidxs] @@ -572,7 +592,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r 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, - mb_count, n_darts_select) + p.size(0), n_darts_select) # mb_count, total_overflow) timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) @@ -591,12 +611,13 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r 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=(mb_count, node_count_total)) + size=(p.size(0), node_count_total)) timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-samplecount") - sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + 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] @@ -631,10 +652,12 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r # 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_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=(mb_count, node_count_total)) + values=next_frontier_values, + size=(p.size(0), node_count_total)) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") @@ -651,9 +674,10 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, r torch.cuda.nvtx.range_push("nvtx-compute-bool") next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] - sampled_count = torch.cuda.IntTensor(mb_count).fill_(0) + 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_nnzidxs, \ - next_frontier._values()[next_frontier_nnzvals].int()) + next_frontier._values()[next_frontier_nnzvals].int()) underfull_minibatches = (sampled_count < frontier_size).any() torch.cuda.nvtx.range_pop() timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index f09d0ef..292eb24 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -17,7 +17,7 @@ from cagnet.nn.conv import GCNConv from cagnet.partitionings import Partitioning -from cagnet.samplers import ladies_sampler +from cagnet.samplers import ladies_sampler, sage_sampler import cagnet.nn.functional as CAGF import socket @@ -367,6 +367,7 @@ def main(args, batches=None): batches_values = torch.cuda.DoubleTensor(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) + print(f"g_loc[3]: {(g_loc._indices()[0, :] == 3).nonzero().squeeze().size()}") torch.cuda.nvtx.range_pop() print(f"g_loc: {g_loc}") @@ -375,7 +376,10 @@ def main(args, batches=None): if args.n_darts == -1: edge_count = adj_matrix.size(1) avg_degree = int(edge_count / node_count) - args.n_darts = avg_degree * args.batch_size + if args.sample_method == "ladies": + args.n_darts = avg_degree * args.batch_size + elif args.sample_method == "sage": + args.n_darts = avg_degree print(f"n_darts: {args.n_darts}") print(f"rank: {rank} g_loc: {g_loc}") @@ -387,24 +391,46 @@ def main(args, batches=None): nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) - torch.manual_seed(rank_col) - current_frontier, next_frontier, adj_matrices = \ - ladies_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) - print() - torch.cuda.profiler.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("nvtx-sampler") - torch.manual_seed(rank_col) - 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.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + if args.sample_method == "ladies": + torch.manual_seed(rank_col) + current_frontier, next_frontier, adj_matrices = \ + ladies_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) + + print() + torch.cuda.profiler.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("nvtx-sampler") + torch.manual_seed(rank_col) + 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.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) + elif args.sample_method == "sage": + torch.manual_seed(rank_col) + current_frontier, next_frontier, adj_matrices = \ + sage_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) + + print() + torch.cuda.profiler.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("nvtx-sampler") + torch.manual_seed(rank_col) + current_frontier, next_frontier, adj_matrices_bulk = \ + sage_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, sa_recv_buff, \ + rank, size, row_groups, col_groups) + torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() @@ -425,9 +451,14 @@ def main(args, batches=None): adj_matrix_sample_indices[0,:] -= row_select_min adj_matrix_sample_values = sampled_values[sample_select_mask] - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) adj_matrices[j][i] = adj_matrix_sample # print(f"sample: {adj_matrices}") @@ -518,6 +549,8 @@ def main(args, batches=None): 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=-1, From 35e2262ceaa8ce1704131839134611ff6694afc7 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 29 Sep 2022 05:26:23 -0400 Subject: [PATCH 099/178] Initial SAGE implementation --- cagnet/samplers/ladies.py | 4 +- cagnet/samplers/sage.py | 41 ++++++++------ cagnet/samplers/utils.py | 77 ++++++++++++++++++++------- sparse-extension/sparse_coo_tensor.cu | 16 +++--- 4 files changed, 95 insertions(+), 43 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 9d2e82f..120cf33 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -52,7 +52,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota nnz = current_frontier[0, :].size(0) p = gen_prob_dist(batches, adj_matrix, mb_count, node_count_total, replication, rank, size, \ - row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict) + row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies") next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ row_groups, col_groups, timing_dict) @@ -60,7 +60,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota batches_select, next_frontier_select, adj_matrix_sample = \ select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ batch_size, frontier_size, mb_count, mb_count_total, node_count_total, \ - replication, rank, size, row_groups, col_groups, timing_dict, i) + replication, rank, size, row_groups, col_groups, timing_dict, i, "ladies") adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 3be91d3..eade8b8 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -25,6 +25,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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) current_frontier = torch.cuda.IntTensor(mb_count_total, batch_size + frontier_size) @@ -71,31 +74,37 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, timing_dict) - # Collapse next_frontier matrix + # add explicit 0's to next_frontier next_frontier_nnz = next_frontier._values().nonzero().squeeze() - frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], - bins=p.size(0)).int() + frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) frontier_nnz_sizes = torch.clamp(frontier_nnz_sizes, max=frontier_size) - print(f"frontier_nnz_sizes: {frontier_nnz_sizes}") - frontier_nnz_sizes = frontier_nnz_sizes.sum() # just for mb=1 next_frontier_rows = torch.repeat_interleave( - torch.arange(mb_count, device=torch.device("cuda:0")), - frontier_nnz_sizes) - next_frontier_idxs = torch.stack( - (next_frontier_rows, - next_frontier._indices()[1, next_frontier_nnz])) + torch.arange(batch_size * mb_count, device=torch.device("cuda:0")), + frontier_size) + nextf_cols_idxs = torch.arange(next_frontier_nnz.size(0), device=torch.device("cuda:0")) + 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_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 next_frontier = torch.sparse_coo_tensor(next_frontier_idxs, - next_frontier._values()[next_frontier_nnz], - size=(mb_count, node_count_total)) - next_frontier = next_frontier.coalesce() - next_frontier._values().fill_(1) + next_frontier_values, + size=(batch_size * mb_count, node_count_total)) + print(f"next_frontier: {next_frontier}") + # next_frontier = next_frontier.coalesce() + # next_frontier._values().fill_(1) batches_select, next_frontier_select, adj_matrix_sample = \ select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, - nnz, batch_size, batch_size * frontier_size, mb_count, + nnz, batch_size, frontier_size, mb_count, mb_count_total, node_count_total, replication, rank, - size, row_groups, col_groups, timing_dict, i) + size, row_groups, col_groups, timing_dict, i, "sage") adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 24942a3..9d5ea58 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -197,9 +197,15 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_cols_count_list = [] for j in range(size // replication): nnz_cols_count_list.append(torch.cuda.IntTensor(1).fill_(0)) - dist.gather(nnz_cols_count, nnz_cols_count_list, dst=q, group=col_groups[rank_col]) + recv_rank = rank_col + j * replication + if recv_rank == q: + nnz_cols_count_list[-1] = nnz_cols_count + else: + dist.irecv(nnz_cols_count[-1], src=recv_rank, group=col_groups[rank_col]).wait() + # 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]) + # dist.gather(nnz_cols_count, dst=q, group=col_groups[rank_col]) + dist.isend(nnz_cols_count, dst=q, group=col_groups[rank_col]) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-gather-nnzcounts-{name}") # Rank q allocates recv buffer for nnz col ids @@ -217,7 +223,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups for j in range(size // replication): recv_rank = rank_col + j * replication if recv_rank != q: - recv_objs.append(dist.irecv(nnz_col_ids[j], src=recv_rank)) + recv_objs.append(dist.irecv(nnz_col_ids[j], src=recv_rank).wait()) else: nnz_col_ids[j] = nnz_cols else: @@ -403,12 +409,16 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) + # p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) + p_den = torch.cuda.FloatTensor(numerator.size(0)).fill_(0) if name == "ladies": p_num_values = torch.square(p_num_values) elif name == "sage": p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) + p_num_values = p_num_values.float() # just for Summit, no 64-bit fp atomicadd scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + p_num_values = p_num_values.double() + p_den = p_den.double() p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(numerator.size(0), node_count_total)) @@ -695,17 +705,22 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ batch_size, frontier_size, mb_count, mb_count_total, node_count_total, replication, \ - rank, size, row_groups, col_groups, timing_dict, layer_id): + rank, size, row_groups, col_groups, timing_dict, layer_id, name): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-construct-nextf") - 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) + 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() timing_dict["construct-nextf"].append(stop_time(start_timer, stop_timer)) @@ -725,10 +740,19 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ row_select_mtx_values = torch.cuda.DoubleTensor(row_select_mtx_indices[0].size(0)).fill_(1.0) repeated_next_frontier = next_frontier_select.clone() - scale_mtx = torch.arange(start=0, end=node_count_total * mb_count, step=node_count_total).cuda() - repeated_next_frontier.add_(scale_mtx[:, None]) + 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() + print(f"scale_mtx: {scale_mtx}") + print(f"next_frontier_select: {next_frontier_select}") + scale_mtx = scale_mtx[:, None] + repeated_next_frontier.add_(scale_mtx) col_select_mtx_rows = repeated_next_frontier.view(-1) - col_select_mtx_cols = torch.arange(next_frontier_select.size(1)).cuda().repeat(mb_count,1).view(-1) + col_select_mtx_cols = torch.arange(next_frontier_select.size(1) * batch_size).cuda().repeat( + mb_count,1).view(-1) + print(f"col_select_mtx_cols: {col_select_mtx_cols}") 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) @@ -747,6 +771,7 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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)) + print(f"row_select_mtx: {row_select_mtx}") sa_masks.fill_(0) sa_recv_buff.fill_(0) sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ @@ -757,13 +782,19 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ torch.cuda.nvtx.range_pop() timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) + print(f"before sample_mtx: {sample_mtx}") # 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) - shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - rank, size, replication, batch_size, node_count_total, mb_count_total) + if name == "ladies": + 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) + elif name == "sage": + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, replication, batch_size, node_count_total, mb_count_total, 1) sampled_indices[1,:] += row_shift torch.cuda.nvtx.range_pop() timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) @@ -771,10 +802,20 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ # 4. Multiply sampled rows with col_select matrix start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-col-select-spgemm") - sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, - size=(nnz * mb_count, node_count_total * mb_count_total)) - 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))) + if name == "ladies": + sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, + size=(nnz * mb_count, node_count_total * mb_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)) + print(f"after sample_mtx: {sample_mtx}") + 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))) + 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)) + print(f"col_select_mtx: {col_select_mtx}") sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ row_groups, col_groups, "colsel", timing_dict) torch.cuda.nvtx.range_pop() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index f45a2df..3f028a0 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -587,7 +587,8 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: } __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 replication, int nnz, int batch_size, int node_count, int mb_count, + int shift_size) { long id = blockIdx.x * blockDim.x + threadIdx.x; long stride = blockDim.x * gridDim.x; @@ -596,14 +597,14 @@ __global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int rank, 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) / batch_size; + long mb_row = (row_select_rows[i] + proc_row_chunk) / shift_size; row_shift[i] += mb_row * 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 mb_count, int shift_size) { int BLOCK_SIZE = 256; @@ -618,7 +619,8 @@ void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_sele nnz, batch_size, node_count, - mb_count); + mb_count, + shift_size); CHECK_ERROR("shift row select error") } @@ -643,7 +645,7 @@ void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, i CHECK_ERROR("shift col select error") } -__global__ void ScatterAddD(double *src, long *indices, double *values, int num_vals) { +__global__ void ScatterAddD(float *src, long *indices, float *values, int num_vals) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -659,9 +661,9 @@ void scatterd_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at int BLOCK_COUNT = std::ceil(num_vals / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ScatterAddD<<>>(src.data(), + ScatterAddD<<>>(src.data(), indices.data(), - values.data(), + values.data(), num_vals); CHECK_ERROR("scatter add doubles error") } From cca9eff4de996a0bd14db639dc4dd0f25290ea60 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 29 Sep 2022 14:38:23 -0700 Subject: [PATCH 100/178] Verified SAGE with test_sage script --- cagnet/samplers/ladies.py | 2 +- cagnet/samplers/sage.py | 5 +- cagnet/samplers/utils.py | 56 +++++------ examples/gcn_15d.py | 6 +- sparse-extension/sparse_coo_tensor.cu | 6 +- tests/test_ladies.py | 56 ----------- tests/test_sage.py | 128 ++++++++++++++++++++++++++ 7 files changed, 160 insertions(+), 99 deletions(-) create mode 100644 tests/test_sage.py diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 120cf33..38bffbb 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -55,7 +55,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies") next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ - row_groups, col_groups, timing_dict) + row_groups, col_groups, timing_dict, "ladies") batches_select, next_frontier_select, adj_matrix_sample = \ select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index eade8b8..ac61c55 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -72,7 +72,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, - timing_dict) + timing_dict, "sage") # add explicit 0's to next_frontier next_frontier_nnz = next_frontier._values().nonzero().squeeze() @@ -88,7 +88,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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_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) @@ -96,7 +96,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier = torch.sparse_coo_tensor(next_frontier_idxs, next_frontier_values, size=(batch_size * mb_count, node_count_total)) - print(f"next_frontier: {next_frontier}") # next_frontier = next_frontier.coalesce() # next_frontier._values().fill_(1) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 9d5ea58..2febebe 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -198,14 +198,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups for j in range(size // replication): nnz_cols_count_list.append(torch.cuda.IntTensor(1).fill_(0)) recv_rank = rank_col + j * replication - if recv_rank == q: - nnz_cols_count_list[-1] = nnz_cols_count - else: - dist.irecv(nnz_cols_count[-1], src=recv_rank, group=col_groups[rank_col]).wait() - # dist.gather(nnz_cols_count, nnz_cols_count_list, dst=q, group=col_groups[rank_col]) + 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]) - dist.isend(nnz_cols_count, dst=q, group=col_groups[rank_col]) + 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}") # Rank q allocates recv buffer for nnz col ids @@ -409,16 +404,12 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-compute-p") - # p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) - p_den = torch.cuda.FloatTensor(numerator.size(0)).fill_(0) + p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) if name == "ladies": p_num_values = torch.square(p_num_values) elif name == "sage": p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) - p_num_values = p_num_values.float() # just for Summit, no 64-bit fp atomicadd scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) - p_num_values = p_num_values.double() - p_den = p_den.double() p = torch.sparse_coo_tensor(indices=p_num_indices, values=p_num_values, size=(numerator.size(0), node_count_total)) @@ -428,7 +419,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication return p def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, - rank, size, row_groups, col_groups, timing_dict): + 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) @@ -474,21 +465,24 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) while underfull_minibatches: iter_count += 1 - print(f"iter_count: {iter_count}", flush=True) - print(f"sampled_count: {sampled_count}") start_time(sample_start_timer) torch.cuda.nvtx.range_push("nvtx-sampling-iter") start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-prob-rowsum") - # ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) - ps_p_values = torch.cumsum(p._values(), dim=0) - # 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 = torch.cumsum(p_rowsum, dim=0) - # ps_p_rowsum[0] = 0 + + 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()) + ps_p_rowsum = torch.cumsum(p_rowsum, dim=0) torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) @@ -559,7 +553,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-psoverflow") selection_iter_count += 1 - print(f"selection_iter_count: {selection_iter_count}", flush=True) ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() # n_darts_select = total_overflow // replication @@ -745,14 +738,15 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ elif name == "sage": scale_mtx = torch.arange(start=0, end=node_count_total * mb_count * batch_size, step=node_count_total).cuda() - print(f"scale_mtx: {scale_mtx}") - print(f"next_frontier_select: {next_frontier_select}") scale_mtx = scale_mtx[:, None] repeated_next_frontier.add_(scale_mtx) col_select_mtx_rows = repeated_next_frontier.view(-1) - col_select_mtx_cols = torch.arange(next_frontier_select.size(1) * batch_size).cuda().repeat( - mb_count,1).view(-1) - print(f"col_select_mtx_cols: {col_select_mtx_cols}") + 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) @@ -771,7 +765,6 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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)) - print(f"row_select_mtx: {row_select_mtx}") sa_masks.fill_(0) sa_recv_buff.fill_(0) sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ @@ -782,7 +775,6 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ torch.cuda.nvtx.range_pop() timing_dict["row-select-spgemm"].append(stop_time(start_timer, stop_timer)) - print(f"before sample_mtx: {sample_mtx}") # 3. Expand sampled rows start_time(start_timer) @@ -808,14 +800,12 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ elif name == "sage": sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, size=(nnz * mb_count, node_count_total * mb_count_total * nnz)) - print(f"after sample_mtx: {sample_mtx}") 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))) 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)) - print(f"col_select_mtx: {col_select_mtx}") sampled_indices, sampled_values = dist_spgemm15D(sample_mtx, col_select_mtx, replication, rank, size, \ row_groups, col_groups, "colsel", timing_dict) torch.cuda.nvtx.range_pop() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 292eb24..2c7ffe9 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -366,7 +366,7 @@ def main(args, batches=None): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) batches_values = torch.cuda.DoubleTensor(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) + g_loc = torch.pow(g_loc, 2) print(f"g_loc[3]: {(g_loc._indices()[0, :] == 3).nonzero().squeeze().size()}") torch.cuda.nvtx.range_pop() print(f"g_loc: {g_loc}") @@ -467,8 +467,8 @@ def main(args, batches=None): torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0), size=(node_count, node_count)) adj_matrix = adj_matrix.coalesce().cuda() - # adj_matrix = row_normalize(adj_matrix) - # adj_matrix = torch.pow(adj_matrix, 2) + adj_matrix = row_normalize(adj_matrix) + adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code return current_frontier, next_frontier, adj_matrices, adj_matrix diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 3f028a0..481d13b 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -645,7 +645,7 @@ void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, i CHECK_ERROR("shift col select error") } -__global__ void ScatterAddD(float *src, long *indices, float *values, int num_vals) { +__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; @@ -661,9 +661,9 @@ void scatterd_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at int BLOCK_COUNT = std::ceil(num_vals / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - ScatterAddD<<>>(src.data(), + ScatterAddD<<>>(src.data(), indices.data(), - values.data(), + values.data(), num_vals); CHECK_ERROR("scatter add doubles error") } diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 795f931..9ded29b 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -10,62 +10,6 @@ eps = 1e-4 def main(args): - """ - # Testing sampled neighbors, probability distribution, and normalized Laplacian - # with original LADIES code - - mine, mine_mtx, mine_batch = examples.gcn_15d.main(args) - mine = mine.cuda() - them, them_mtx, them_batch = LADIES.pytorch_ladies.main(args) - them = torch.from_numpy(them).cuda() - - mine_sampled_neighbors = mine._indices()[1,:].squeeze() - them_sampled_neighbors = np.nonzero(them).squeeze() - - mine_sampled_probs = mine._values().float() - them_sampled_probs = them[them_sampled_neighbors].float() - - if not torch.equal(mine_sampled_neighbors, them_sampled_neighbors): - print("indices don't match") - print(f"mine_sampled_neighbors.size: {mine_sampled_neighbors.size()}") - print(f"mine_sampled_neighbors: {mine_sampled_neighbors}") - print(f"them_sampled_neighbors.size: {them_sampled_neighbors.size()}") - print(f"them_sampled_neighbors: {them_sampled_neighbors}") - else: - print("indices match") - - if not torch.allclose(mine_sampled_probs, them_sampled_probs, atol=0.001, rtol=0.0): - print("probabilities don't match") - print(f"mine_sampled_probs.size: {mine_sampled_probs.size()}") - print(f"mine_sampled_probs: {mine_sampled_probs}") - print(f"them_sampled_probs.size: {them_sampled_probs.size()}") - print(f"them_sampled_probs: {them_sampled_probs}") - else: - print("probabilities match") - - mine_mtx_normalization = mine_mtx._values() - them_mtx_normalization = torch.from_numpy(them_mtx.data).float().cuda() - - if not torch.allclose(mine_mtx_normalization, them_mtx_normalization, atol=0.001, rtol=0.0): - print("laplacians don't match") - unmatch_idx = ~torch.isclose(mine_mtx_normalization, them_mtx_normalization, atol=0.001, rtol=0.0) - print(f"unmatch_idx: {unmatch_idx.nonzero()}") - print(f"mine_mtx_normalization: {mine_mtx_normalization[unmatch_idx]}") - print(f"them_mtx_normalization: {them_mtx_normalization[unmatch_idx]}") - else: - print("laplacians match") - - them_batch = torch.from_numpy(them_batch).cuda() - - if not torch.equal(mine_batch, them_batch): - print("batches don't match") - print(f"mine_batch: {mine_batch}") - print(f"them_batch: {them_batch}") - else: - print("batches match") - """ - - """ All vertices in new frontier exist in aggregated neighborhood of old frontier """ diff --git a/tests/test_sage.py b/tests/test_sage.py new file mode 100644 index 0000000..bf7c959 --- /dev/null +++ b/tests/test_sage.py @@ -0,0 +1,128 @@ +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 SAGE sampler test") + 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 {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 + + mb_id = i // args.batch_size + adj_matrices = adj_matrices_all[mb_id] # adj_matrices for minibatch i + + vtx_mask = adj_matrices[0]._indices()[0,:] == (i % args.batch_size) + vtx_indices = adj_matrices[0]._indices()[:, vtx_mask] + vtx_values = adj_matrices[0]._values()[vtx_mask] + adj_matrices = [torch.sparse_coo_tensor(vtx_indices, vtx_values, size=adj_matrices[0].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: + 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("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='sage', + help='Sampled Algorithms: ladies/fastgcn/full') + parser.add_argument('--batch-num', type=int, default= 10, + help='Maximum Batch Number') + parser.add_argument('--pool-num', type=int, default= 10, + help='Number of Pool') + parser.add_argument('--n-bulkmb', type=int, default=1, + help='Number of minibatches to sample in bulk') + parser.add_argument('--n-darts', type=int, default=-1, + help='Number of darts to throw per minibatch in LADIES sampling') + args = parser.parse_args() + print(args) + + main(args) From 9873dbee50e7aadee587aa91f471b7ebdd468b44 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 1 Oct 2022 14:27:41 -0700 Subject: [PATCH 101/178] Added erroneous spgemm extension --- .gitignore | 2 + cagnet/samplers/sage.py | 29 ++- cagnet/samplers/utils.py | 160 ++++++++++----- examples/gcn_15d.py | 12 +- sparse-extension/sparse_coo_tensor.cu | 281 +++++++++++++++++++++++--- 5 files changed, 396 insertions(+), 88 deletions(-) diff --git a/.gitignore b/.gitignore index 15f95b3..3cc0263 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,8 @@ reddit.py tests/data tests/examples tests/LADIES/ +tests/pytorch_geometric/ +tests/torch-quiver/ tests/*.out *.qdrep *.out diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index ac61c55..72b0582 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -9,15 +9,23 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu + scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ + sparse_coo_tensor_gpu + + +timing = True def start_time(timer): - timer.record() + if timing == True: + timer.record() def stop_time(start_timer, stop_timer): - stop_timer.record() - torch.cuda.synchronize() - return start_timer.elapsed_time(stop_timer) + if timing == True: + stop_timer.record() + torch.cuda.synchronize() + return start_timer.elapsed_time(stop_timer) + else: + return 0.0 def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups): @@ -48,7 +56,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # adj_matrices[i][j] -- mb i layer j adj_matrices = [None] * n_layers # adj_matrices[i] -- bulk mb mtx for layer j - start_time(total_start_timer) + # start_time(total_start_timer) for i in range(n_layers): if i == 0: nnz = batch_size @@ -56,6 +64,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, nnz = current_frontier[0, :].size(0) # Expand batches matrix + total_start_timer.record() batches_expand_rows = torch.arange(mb_count * nnz, \ device=torch.device("cuda:0")) batches_expand_idxs = torch.stack( @@ -73,6 +82,10 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, timing_dict, "sage") + total_stop_timer.record() + torch.cuda.synchronize() + total_time = total_start_timer.elapsed_time(total_stop_timer) + print(f"total_time: {total_time}", flush=True) # add explicit 0's to next_frontier next_frontier_nnz = next_frontier._values().nonzero().squeeze() @@ -93,9 +106,11 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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 + next_frontier = torch.sparse_coo_tensor(next_frontier_idxs, next_frontier_values, size=(batch_size * mb_count, node_count_total)) + # next_frontier = next_frontier.coalesce() # next_frontier._values().fill_(1) @@ -108,8 +123,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier - - print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) 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": # print(f"{k} times: {v}") diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 2febebe..d441b47 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -8,15 +8,23 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu + scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ + sparse_coo_tensor_gpu, spgemm_gpu + + +timing = True def start_time(timer): - timer.record() + if timing == True: + timer.record() def stop_time(start_timer, stop_timer): - stop_timer.record() - torch.cuda.synchronize() - return start_timer.elapsed_time(stop_timer) + if timing == True: + stop_timer.record() + torch.cuda.synchronize() + return start_timer.elapsed_time(stop_timer) + else: + return 0.0 def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if timing_dict is not None: @@ -81,14 +89,44 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}") start_time(start_timer) - 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_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) + + print(f"before spgemm_gpu", flush=True) + print(f"mata_chunk_indices[0,:]: {mata_chunk_indices[0,:].int()}", flush=True) + print(f"mata_chunk_indices[0,:].size: {mata_chunk_indices[0,:].int().size()}", flush=True) + print(f"mata.size(0): {mata.size(0)}", flush=True) + mata_chunk_indices = mata_chunk_indices.int() + matb_recv_indices = matb_recv_indices.int() + matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows + mata_chunk_indices[1,:], # A_cols + mata_chunk_values, # A_vals + matb_recv_indices[0,:], # B_rows + matb_recv_indices[1,:], # B_cols + matb_recv_values, # B_vals + mata.size(0), # n + chunk_col_size, # k + matb.size(1)) # m + print(f"after spgemm_gpu", flush=True) + mata_chunk_indices = mata_chunk_indices.long() + matb_recv_indices = matb_recv_indices.long() + matc_chunk_rows = matc_outputs[0] + matc_chunk_cols = matc_outputs[1] + matc_chunk_values = matc_outputs[2] + + print(f"before matc_chunk_rows: {matc_chunk_rows}") + matc_chunk_counts = torch.diff(matc_chunk_rows) + matc_chunk_rows = torch.repeat_interleave(torch.arange(0, mata.size(0), device=torch.device("cuda:0")), + matc_chunk_counts) + print(f"after matc_chunk_rows: {matc_chunk_rows}") + matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") start_time(start_timer) - matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + # 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}") start_time(start_timer) @@ -130,7 +168,8 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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(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}") start_time(start_timer) @@ -315,6 +354,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_values = matb_select_recv[2, :].double() start_time(start_timer) + print(f"type(mata_chunk_indices): {mata_chunk_indices.dtype}") matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ matb_recv_values, mata.size(0), \ @@ -322,7 +362,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") start_time(start_timer) - matc_chunk = torch.sparse_coo_tensor(matc_chunk_indices, matc_chunk_values, size=matc.size()) + # 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, + torch.Size([matc.size(0), matc.size(1)])) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") start_time(start_timer) @@ -364,7 +406,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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(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}") start_time(start_timer) @@ -410,9 +453,10 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication elif name == "sage": p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) - p = torch.sparse_coo_tensor(indices=p_num_indices, - values=p_num_values, - size=(numerator.size(0), node_count_total)) + # 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])) print(f"p.nnz: {p._nnz()}") normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) @@ -443,9 +487,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") - next_frontier = torch.sparse_coo_tensor(indices=p._indices(), - values=torch.cuda.LongTensor(p._nnz()).fill_(0), - size=(p.size(0), node_count_total)) + # next_frontier = torch.sparse_coo_tensor(indices=p._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])) frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], bins=p.size(0)).int() @@ -482,7 +528,9 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, ps_p_values = torch.cumsum(p._values(), dim=0) p_rowsum.fill_(0) p_rowsum.scatter_add_(0, p._indices()[0, :], p._values()) - ps_p_rowsum = torch.cumsum(p_rowsum, dim=0) + # 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 torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) @@ -517,9 +565,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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 = 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])) torch.cuda.nvtx.range_pop() timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) @@ -569,9 +619,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-instmtx") - dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), - values=dart_hits_inv, - size=(p.size(0), node_count_total)) + # dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), + # values=dart_hits_inv, + # size=(p.size(0), node_count_total)) + 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) @@ -612,9 +664,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") # 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 = 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) @@ -658,9 +712,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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 = 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])) start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-set-probs") @@ -763,15 +819,19 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ # 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 = 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])) sa_masks.fill_(0) sa_recv_buff.fill_(0) sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ row_groups, col_groups, "rowsel", sa_masks, sa_recv_buff, timing_dict) - sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, - size=(nnz * mb_count, node_count_total)) + # 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)) @@ -795,17 +855,25 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 = 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 * mb_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 = 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 * mb_count_total * nnz])) 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 = 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)])) 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, - size=(node_count_total * mb_count * batch_size, next_frontier.size(1) * batch_size)) + torch.Size([node_count_total * mb_count * batch_size, next_frontier.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) torch.cuda.nvtx.range_pop() @@ -814,8 +882,10 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ # 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 = 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)])) torch.cuda.nvtx.range_pop() timing_dict["set-sample"].append(stop_time(start_timer, stop_timer)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 2c7ffe9..498ad30 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -303,6 +303,7 @@ def main(args, batches=None): print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + edge_count = adj_matrix.size(1) partitioning = Partitioning.ONE5D @@ -317,6 +318,8 @@ def main(args, batches=None): features_loc, g_loc, ampbyp = one5d_partition(rank, size, inputs, adj_matrix, data, \ inputs, num_classes, args.replication, \ args.normalize) + del inputs + print("end partitioning", flush=True) features_loc = features_loc.to(device) @@ -367,14 +370,12 @@ def main(args, batches=None): batches_values = torch.cuda.DoubleTensor(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) - print(f"g_loc[3]: {(g_loc._indices()[0, :] == 3).nonzero().squeeze().size()}") torch.cuda.nvtx.range_pop() print(f"g_loc: {g_loc}") - node_count = inputs.size(0) + node_count = g_loc.size(1) adj_matrix = adj_matrix.cuda() if args.n_darts == -1: - edge_count = adj_matrix.size(1) avg_degree = int(edge_count / node_count) if args.sample_method == "ladies": args.n_darts = avg_degree * args.batch_size @@ -382,6 +383,8 @@ def main(args, batches=None): args.n_darts = avg_degree print(f"n_darts: {args.n_darts}") + # del adj_matrix # Comment when testing + print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing @@ -389,7 +392,7 @@ def main(args, batches=None): nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) - nnz_recv_upperbound = adj_matrix.size(1) // (size // args.replication) + nnz_recv_upperbound = edge_count // (size // args.replication) sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) if args.sample_method == "ladies": @@ -462,6 +465,7 @@ def main(args, batches=None): adj_matrices[j][i] = adj_matrix_sample # print(f"sample: {adj_matrices}") + # 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), diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 481d13b..d65b9fe 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -188,10 +188,6 @@ void spmm_gpu(const at::Tensor& A_rowindices, 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 @@ -271,7 +267,8 @@ void spmm_gpu(const at::Tensor& A_rowindices, void* d_buffer = NULL; - CHECK_ERROR(cudaMalloc(&d_buffer, bufferSize)); + // CHECK_ERROR(cudaMalloc(&d_buffer, bufferSize)); + cudaMalloc(&d_buffer, bufferSize); CHECK_CUSPARSE(cusparseSpMM(handle, // handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA @@ -294,32 +291,253 @@ void spmm_gpu(const at::Tensor& A_rowindices, C.t_(); } -// __global__ void DownSample(long *q_values, long *q_rows, int *overflow, int nnz) { -// int id = blockIdx.x * blockDim.x + threadIdx.x; -// int stride = blockDim.x * gridDim.x; -// -// for (int i = id; i < nnz; i += stride) { -// long row = q_rows[i]; -// if (q_values[i] == 1 && ((int) atomicSub((unsigned int *)&overflow[row], 1)) > 0) { -// q_values[i] = 0; -// } -// } -// } - -// void downsample_gpu(const at::Tensor& q_values, -// const at::Tensor& q_rows, -// 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<<>>(q_values.data(), q_rows.data(), -// overflow.data(), nnz); -// CHECK_ERROR("downsampling error") -// } +std::vector spgemm_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 n, + int32_t k, + int32_t m) { + + + printf("before spgemm\n"); fflush(stdout); + // SpGEMM between A (n x k) and B (k x m) + auto handle = at::cuda::getCurrentCUDASparseHandle(); + + int32_t *d_a_csrrows = NULL; + int32_t *d_b_csrrows = NULL; + int *tmp; + + int32_t A_nnz = A_values.size(0); + int32_t B_nnz = B_values.size(0); + printf("A_nnz: %d\n", A_nnz); fflush(stdout); + printf("B_nnz: %d\n", B_nnz); fflush(stdout); + + printf("before coo2csr\n"); fflush(stdout); + // Construct CSR offsets array for A and B + // CHECK_ERROR(cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t))); + cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); + std::cout << "d_a_csrrows: " << d_a_csrrows << std::endl; + cudaMemset(d_a_csrrows, 0, (n + 1) * sizeof(int32_t)); + std::cout << "n: " << n << std::endl; + std::cout << "handle: " << handle << std::endl; + cudaDeviceSynchronize(); + cudaError_t err; + err = cudaGetLastError(); + printf("err: %s\n", cudaGetErrorString(err)); fflush(stdout); + + int *d_a_indices = A_rowindices.data(); + printf("d_a_indices: %x\n", d_a_indices); fflush(stdout); + int *h_a_indices = new int[A_nnz](); + cudaMemcpy(h_a_indices, d_a_indices, A_nnz * sizeof(int), cudaMemcpyDeviceToHost); + for (int i = 0; i < A_nnz; i++) { + printf("%d ", h_a_indices[i]); fflush(stdout); + } + printf("\n"); fflush(stdout); + cudaDeviceSynchronize(); + err = cudaGetLastError(); + printf("after printouts\n"); fflush(stdout); + printf("err: %s\n", cudaGetErrorString(err)); fflush(stdout); + CHECK_CUSPARSE(cusparseXcoo2csr(handle, + A_rowindices.data(), + A_nnz, + n, + d_a_csrrows, + CUSPARSE_INDEX_BASE_ZERO)); + cudaDeviceSynchronize(); + err = cudaGetLastError(); + printf("after coo2csr\n"); fflush(stdout); + printf("err: %s\n", cudaGetErrorString(err)); fflush(stdout); + // CHECK_ERROR(cudaMalloc(&d_b_csrrows, (k + 1) * sizeof(int32_t))); + cudaMalloc(&d_b_csrrows, (k + 1) * sizeof(int32_t)); + printf("after balloc\n"); fflush(stdout); + + CHECK_CUSPARSE(cusparseXcoo2csr(handle, + B_rowindices.data(), + B_nnz, + k, + d_b_csrrows, + CUSPARSE_INDEX_BASE_ZERO)); + printf("after coo2csr\n"); fflush(stdout); + + float alpha = 1; + float beta = 0; + + // Construct CSR matrix structs for A, B, and C (empty CSR) + printf("before csr_construct\n"); fflush(stdout); + cusparseSpMatDescr_t matA; + CHECK_CUSPARSE(cusparseCreateCsr(&matA, + n, // rows + k, // cols + A_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_64F)); // valueType + + cusparseSpMatDescr_t matB; + CHECK_CUSPARSE(cusparseCreateCsr(&matB, + k, // rows + m, // cols + B_nnz, // nnz + d_b_csrrows, // csrRowOffsets + B_colindices.data(), // csrColInd + B_values.data(), // csrValues + CUSPARSE_INDEX_32I, // csrRowOffsetsType + CUSPARSE_INDEX_32I, // csrColIndType + CUSPARSE_INDEX_BASE_ZERO, // idxBase, + CUDA_R_64F)); // valueType + + cusparseSpMatDescr_t matC; + CHECK_CUSPARSE(cusparseCreateCsr(&matC, + n, // rows + m, // cols + 0, // nnz + NULL, // csrRowOffsets + NULL, // csrColInd + NULL, // csrValues + CUSPARSE_INDEX_32I, // csrRowOffsetsType + CUSPARSE_INDEX_32I, // csrColIndType + CUSPARSE_INDEX_BASE_ZERO, // idxBase, + CUDA_R_64F)); // valueType + printf("after csr_construct\n"); fflush(stdout); + + printf("before spgemm_workest\n"); fflush(stdout); + cusparseSpGEMMDescr_t spgemmDescr; + CHECK_CUSPARSE(cusparseSpGEMM_createDescr(&spgemmDescr)); + + void *dBuffer1 = NULL; + void *dBuffer2 = NULL; + size_t bufferSize1 = 0; + size_t bufferSize2 = 0; + printf("bufferSize1: %ld\n", bufferSize1); fflush(stdout); + 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_64F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize1, // bufferSize1 + NULL)); // externalBuffer1 + printf("after spgemm_workest1\n"); fflush(stdout); + + printf("sizeof(bufferSize1): %d\n", sizeof(bufferSize1)); fflush(stdout); + printf("bufferSize1: %ld\n", bufferSize1); fflush(stdout); + //CHECK_ERROR(cudaMalloc(&dBuffer1, bufferSize1)); + CHECK_ERROR("before cudamalloc error"); + cudaMalloc(&dBuffer1, bufferSize1); + CHECK_ERROR("after cudamalloc error"); + printf("before cudamalloc1 %x\n", dBuffer1); fflush(stdout); + printf("after cudamalloc1 %x\n", dBuffer1); fflush(stdout); + + 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_64F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize1, // bufferSize1 + dBuffer1)); // externalBuffer1 + printf("after workest2\n"); fflush(stdout); + + 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_64F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize2, // bufferSize2 + NULL)); // externalBuffer2 + printf("after spgemm1\n"); fflush(stdout); + + // CHECK_ERROR(cudaMalloc((void**)&dBuffer2, bufferSize2)); + cudaMalloc((void**)&dBuffer2, bufferSize2); + printf("after cudamalloc2\n"); fflush(stdout); + + 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_64F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr, // spgemmDescr + &bufferSize2, // bufferSize2 + dBuffer2)); // externalBuffer2 + printf("after spgemm2\n"); fflush(stdout); + + 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; + double *dC_values; + cudaMalloc(&dC_csrOffsets, (n + 1) * sizeof(int32_t)); + cudaMalloc(&dC_columns, C_nnz1 * sizeof(int32_t)); + cudaMalloc(&dC_values, C_nnz1 * sizeof(double)); + printf("after cudamalloc3\n"); fflush(stdout); + + 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_64F, // computeType + CUSPARSE_SPGEMM_DEFAULT, // alg + spgemmDescr)); + printf("after spgemm_copy\n"); fflush(stdout); + + + cudaFree(dBuffer1); + cudaFree(dBuffer2); + + CHECK_CUSPARSE(cusparseSpGEMM_destroyDescr(spgemmDescr)); + CHECK_CUSPARSE(cusparseDestroySpMat(matA)); + CHECK_CUSPARSE(cusparseDestroySpMat(matB)); + CHECK_CUSPARSE(cusparseDestroySpMat(matC)); + printf("after destroys\n"); fflush(stdout); + + auto rows_dim = torch::IntArrayRef{n}; + auto nnz_dim = torch::IntArrayRef{C_nnz1}; + + auto c_csrOffsets = torch::from_blob(dC_csrOffsets, rows_dim, torch::TensorOptions().device(torch::kCUDA)); + auto c_columns = torch::from_blob(dC_columns, nnz_dim, torch::TensorOptions().device(torch::kCUDA)); + auto c_values = torch::from_blob(dC_values, nnz_dim, torch::TensorOptions().device(torch::kCUDA)); + printf("after tensor_casts\n"); fflush(stdout); + + CHECK_ERROR("error"); + printf("after spgemm\n"); fflush(stdout); + 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) { @@ -726,6 +944,7 @@ void rowselect_coo_gpu(const at::Tensor& nnz_cols, const at::Tensor& rows, const 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"); + 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"); From cb2c31b190ec0d08122c6944cf7b4d8ffd0e4e42 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 1 Oct 2022 18:06:53 -0700 Subject: [PATCH 102/178] spgemm_gpu() extension --- cagnet/samplers/utils.py | 94 ++++++++++++++--------- examples/test_spgemm.py | 6 ++ sparse-extension/sparse_coo_tensor.cu | 106 +++++++------------------- 3 files changed, 91 insertions(+), 115 deletions(-) create mode 100644 examples/test_spgemm.py diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index d441b47..b1cf433 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -89,39 +89,39 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}") start_time(start_timer) - # 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) - - print(f"before spgemm_gpu", flush=True) - print(f"mata_chunk_indices[0,:]: {mata_chunk_indices[0,:].int()}", flush=True) - print(f"mata_chunk_indices[0,:].size: {mata_chunk_indices[0,:].int().size()}", flush=True) - print(f"mata.size(0): {mata.size(0)}", flush=True) - mata_chunk_indices = mata_chunk_indices.int() - matb_recv_indices = matb_recv_indices.int() - matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows - mata_chunk_indices[1,:], # A_cols - mata_chunk_values, # A_vals - matb_recv_indices[0,:], # B_rows - matb_recv_indices[1,:], # B_cols - matb_recv_values, # B_vals - mata.size(0), # n - chunk_col_size, # k - matb.size(1)) # m - print(f"after spgemm_gpu", flush=True) - mata_chunk_indices = mata_chunk_indices.long() - matb_recv_indices = matb_recv_indices.long() - matc_chunk_rows = matc_outputs[0] - matc_chunk_cols = matc_outputs[1] - matc_chunk_values = matc_outputs[2] + 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) - print(f"before matc_chunk_rows: {matc_chunk_rows}") - matc_chunk_counts = torch.diff(matc_chunk_rows) - matc_chunk_rows = torch.repeat_interleave(torch.arange(0, mata.size(0), device=torch.device("cuda:0")), - matc_chunk_counts) - print(f"after matc_chunk_rows: {matc_chunk_rows}") - matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + # print(f"before spgemm_gpu", flush=True) + # print(f"mata_chunk_indices[0,:]: {mata_chunk_indices[0,:].int()}", flush=True) + # print(f"mata_chunk_indices[0,:].size: {mata_chunk_indices[0,:].int().size()}", flush=True) + # print(f"mata.size(0): {mata.size(0)}", flush=True) + # mata_chunk_indices = mata_chunk_indices.int() + # matb_recv_indices = matb_recv_indices.int() + # matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows + # mata_chunk_indices[1,:], # A_cols + # mata_chunk_values.float(), # A_vals + # matb_recv_indices[0,:], # B_rows + # matb_recv_indices[1,:], # B_cols + # matb_recv_values.float(), # B_vals + # mata.size(0), # n + # chunk_col_size, # k + # matb.size(1)) # m + # print(f"after spgemm_gpu", flush=True) + # mata_chunk_indices = mata_chunk_indices.long() + # matb_recv_indices = matb_recv_indices.long() + # matc_chunk_rows = matc_outputs[0].long() + # matc_chunk_cols = matc_outputs[1].long() + # matc_chunk_values = matc_outputs[2].double() + + # print(f"before matc_chunk_rows: {matc_chunk_rows}") + # matc_chunk_counts = torch.diff(matc_chunk_rows) + # matc_chunk_rows = torch.repeat_interleave(torch.arange(0, mata.size(0), device=torch.device("cuda:0")), + # matc_chunk_counts) + # print(f"after matc_chunk_rows: {matc_chunk_rows}") + # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") start_time(start_timer) @@ -354,11 +354,31 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_values = matb_select_recv[2, :].double() start_time(start_timer) - print(f"type(mata_chunk_indices): {mata_chunk_indices.dtype}") - 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_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) + mata_chunk_indices = mata_chunk_indices.int() + matb_recv_indices = matb_recv_indices.int() + matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows + mata_chunk_indices[1,:], # A_cols + mata_chunk_values.float(), # A_vals + matb_recv_indices[0,:], # B_rows + matb_recv_indices[1,:], # B_cols + matb_recv_values.float(), # B_vals + mata.size(0), # n + chunk_col_size, # k + matb.size(1)) # m + mata_chunk_indices = mata_chunk_indices.long() + matb_recv_indices = matb_recv_indices.long() + 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, mata.size(0), device=torch.device("cuda:0")), + matc_chunk_counts) + matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") start_time(start_timer) 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/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index d65b9fe..e18e4fc 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -303,71 +303,36 @@ std::vector spgemm_gpu( int32_t m) { - printf("before spgemm\n"); fflush(stdout); // SpGEMM between A (n x k) and B (k x m) auto handle = at::cuda::getCurrentCUDASparseHandle(); int32_t *d_a_csrrows = NULL; int32_t *d_b_csrrows = NULL; - int *tmp; int32_t A_nnz = A_values.size(0); int32_t B_nnz = B_values.size(0); - printf("A_nnz: %d\n", A_nnz); fflush(stdout); - printf("B_nnz: %d\n", B_nnz); fflush(stdout); - printf("before coo2csr\n"); fflush(stdout); // Construct CSR offsets array for A and B - // CHECK_ERROR(cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t))); cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); - std::cout << "d_a_csrrows: " << d_a_csrrows << std::endl; - cudaMemset(d_a_csrrows, 0, (n + 1) * sizeof(int32_t)); - std::cout << "n: " << n << std::endl; - std::cout << "handle: " << handle << std::endl; - cudaDeviceSynchronize(); - cudaError_t err; - err = cudaGetLastError(); - printf("err: %s\n", cudaGetErrorString(err)); fflush(stdout); - - int *d_a_indices = A_rowindices.data(); - printf("d_a_indices: %x\n", d_a_indices); fflush(stdout); - int *h_a_indices = new int[A_nnz](); - cudaMemcpy(h_a_indices, d_a_indices, A_nnz * sizeof(int), cudaMemcpyDeviceToHost); - for (int i = 0; i < A_nnz; i++) { - printf("%d ", h_a_indices[i]); fflush(stdout); - } - printf("\n"); fflush(stdout); - cudaDeviceSynchronize(); - err = cudaGetLastError(); - printf("after printouts\n"); fflush(stdout); - printf("err: %s\n", cudaGetErrorString(err)); fflush(stdout); CHECK_CUSPARSE(cusparseXcoo2csr(handle, A_rowindices.data(), A_nnz, n, d_a_csrrows, CUSPARSE_INDEX_BASE_ZERO)); - cudaDeviceSynchronize(); - err = cudaGetLastError(); - printf("after coo2csr\n"); fflush(stdout); - printf("err: %s\n", cudaGetErrorString(err)); fflush(stdout); - // CHECK_ERROR(cudaMalloc(&d_b_csrrows, (k + 1) * sizeof(int32_t))); cudaMalloc(&d_b_csrrows, (k + 1) * sizeof(int32_t)); - printf("after balloc\n"); fflush(stdout); - + CHECK_CUSPARSE(cusparseXcoo2csr(handle, B_rowindices.data(), B_nnz, k, d_b_csrrows, CUSPARSE_INDEX_BASE_ZERO)); - printf("after coo2csr\n"); fflush(stdout); float alpha = 1; float beta = 0; // Construct CSR matrix structs for A, B, and C (empty CSR) - printf("before csr_construct\n"); fflush(stdout); cusparseSpMatDescr_t matA; CHECK_CUSPARSE(cusparseCreateCsr(&matA, n, // rows @@ -375,11 +340,11 @@ std::vector spgemm_gpu( A_nnz, // nnz d_a_csrrows, // csrRowOffsets A_colindices.data(), // csrColInd - A_values.data(), // csrValues + A_values.data(), // csrValues CUSPARSE_INDEX_32I, // csrRowOffsetsType CUSPARSE_INDEX_32I, // csrColIndType CUSPARSE_INDEX_BASE_ZERO, // idxBase, - CUDA_R_64F)); // valueType + CUDA_R_32F)); // valueType cusparseSpMatDescr_t matB; CHECK_CUSPARSE(cusparseCreateCsr(&matB, @@ -388,11 +353,11 @@ std::vector spgemm_gpu( B_nnz, // nnz d_b_csrrows, // csrRowOffsets B_colindices.data(), // csrColInd - B_values.data(), // csrValues + B_values.data(), // csrValues CUSPARSE_INDEX_32I, // csrRowOffsetsType CUSPARSE_INDEX_32I, // csrColIndType CUSPARSE_INDEX_BASE_ZERO, // idxBase, - CUDA_R_64F)); // valueType + CUDA_R_32F)); // valueType cusparseSpMatDescr_t matC; CHECK_CUSPARSE(cusparseCreateCsr(&matC, @@ -405,10 +370,8 @@ std::vector spgemm_gpu( CUSPARSE_INDEX_32I, // csrRowOffsetsType CUSPARSE_INDEX_32I, // csrColIndType CUSPARSE_INDEX_BASE_ZERO, // idxBase, - CUDA_R_64F)); // valueType - printf("after csr_construct\n"); fflush(stdout); - - printf("before spgemm_workest\n"); fflush(stdout); + CUDA_R_32F)); // valueType + cusparseSpGEMMDescr_t spgemmDescr; CHECK_CUSPARSE(cusparseSpGEMM_createDescr(&spgemmDescr)); @@ -416,7 +379,6 @@ std::vector spgemm_gpu( void *dBuffer2 = NULL; size_t bufferSize1 = 0; size_t bufferSize2 = 0; - printf("bufferSize1: %ld\n", bufferSize1); fflush(stdout); CHECK_CUSPARSE(cusparseSpGEMM_workEstimation(handle, // handle CUSPARSE_OPERATION_NON_TRANSPOSE, // opA CUSPARSE_OPERATION_NON_TRANSPOSE, // opB @@ -425,21 +387,14 @@ std::vector spgemm_gpu( matB, // matB &beta, // beta matC, // matC - CUDA_R_64F, // computeType + CUDA_R_32F, // computeType CUSPARSE_SPGEMM_DEFAULT, // alg spgemmDescr, // spgemmDescr &bufferSize1, // bufferSize1 NULL)); // externalBuffer1 - printf("after spgemm_workest1\n"); fflush(stdout); - printf("sizeof(bufferSize1): %d\n", sizeof(bufferSize1)); fflush(stdout); - printf("bufferSize1: %ld\n", bufferSize1); fflush(stdout); //CHECK_ERROR(cudaMalloc(&dBuffer1, bufferSize1)); - CHECK_ERROR("before cudamalloc error"); cudaMalloc(&dBuffer1, bufferSize1); - CHECK_ERROR("after cudamalloc error"); - printf("before cudamalloc1 %x\n", dBuffer1); fflush(stdout); - printf("after cudamalloc1 %x\n", dBuffer1); fflush(stdout); CHECK_CUSPARSE(cusparseSpGEMM_workEstimation(handle, // handle CUSPARSE_OPERATION_NON_TRANSPOSE, // opA @@ -449,12 +404,11 @@ std::vector spgemm_gpu( matB, // matB &beta, // beta matC, // matC - CUDA_R_64F, // computeType + CUDA_R_32F, // computeType CUSPARSE_SPGEMM_DEFAULT, // alg spgemmDescr, // spgemmDescr &bufferSize1, // bufferSize1 dBuffer1)); // externalBuffer1 - printf("after workest2\n"); fflush(stdout); CHECK_CUSPARSE(cusparseSpGEMM_compute(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA @@ -464,16 +418,13 @@ std::vector spgemm_gpu( matB, // matB &beta, // beta matC, // matC - CUDA_R_64F, // computeType + CUDA_R_32F, // computeType CUSPARSE_SPGEMM_DEFAULT, // alg spgemmDescr, // spgemmDescr &bufferSize2, // bufferSize2 NULL)); // externalBuffer2 - printf("after spgemm1\n"); fflush(stdout); - // CHECK_ERROR(cudaMalloc((void**)&dBuffer2, bufferSize2)); - cudaMalloc((void**)&dBuffer2, bufferSize2); - printf("after cudamalloc2\n"); fflush(stdout); + cudaMalloc(&dBuffer2, bufferSize2); CHECK_CUSPARSE(cusparseSpGEMM_compute(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, // opA @@ -483,23 +434,21 @@ std::vector spgemm_gpu( matB, // matB &beta, // beta matC, // matC - CUDA_R_64F, // computeType + CUDA_R_32F, // computeType CUSPARSE_SPGEMM_DEFAULT, // alg spgemmDescr, // spgemmDescr &bufferSize2, // bufferSize2 dBuffer2)); // externalBuffer2 - printf("after spgemm2\n"); fflush(stdout); 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; - double *dC_values; + 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(double)); - printf("after cudamalloc3\n"); fflush(stdout); + cudaMalloc(&dC_values, C_nnz1 * sizeof(float)); CHECK_CUSPARSE(cusparseCsrSetPointers(matC, dC_csrOffsets, dC_columns, dC_values)); CHECK_CUSPARSE(cusparseSpGEMM_copy(handle, @@ -510,12 +459,24 @@ std::vector spgemm_gpu( matB, // matB &beta, // beta matC, // matC - CUDA_R_64F, // computeType + CUDA_R_32F, // computeType CUSPARSE_SPGEMM_DEFAULT, // alg spgemmDescr)); - printf("after spgemm_copy\n"); fflush(stdout); + 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); @@ -523,18 +484,7 @@ std::vector spgemm_gpu( CHECK_CUSPARSE(cusparseDestroySpMat(matA)); CHECK_CUSPARSE(cusparseDestroySpMat(matB)); CHECK_CUSPARSE(cusparseDestroySpMat(matC)); - printf("after destroys\n"); fflush(stdout); - - auto rows_dim = torch::IntArrayRef{n}; - auto nnz_dim = torch::IntArrayRef{C_nnz1}; - - auto c_csrOffsets = torch::from_blob(dC_csrOffsets, rows_dim, torch::TensorOptions().device(torch::kCUDA)); - auto c_columns = torch::from_blob(dC_columns, nnz_dim, torch::TensorOptions().device(torch::kCUDA)); - auto c_values = torch::from_blob(dC_values, nnz_dim, torch::TensorOptions().device(torch::kCUDA)); - printf("after tensor_casts\n"); fflush(stdout); - CHECK_ERROR("error"); - printf("after spgemm\n"); fflush(stdout); return {c_csrOffsets, c_columns, c_values}; } From 8bc0ae4ac47b151be28107caca1dd69a7bf5e975 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 3 Oct 2022 22:17:29 -0700 Subject: [PATCH 103/178] Verified Cora/Reddit GraphSAGE accuracy --- cagnet/samplers/sage.py | 23 +++-- cagnet/samplers/utils.py | 121 +++++++++++++++----------- examples/gcn_15d.py | 100 +++++++++++---------- sparse-extension/sparse_coo_tensor.cu | 1 + 4 files changed, 145 insertions(+), 100 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 72b0582..9c8cbaf 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -14,6 +14,7 @@ timing = True +baseline_compare = False def start_time(timer): if timing == True: @@ -56,7 +57,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # adj_matrices[i][j] -- mb i layer j adj_matrices = [None] * n_layers # adj_matrices[i] -- bulk mb mtx for layer j - # start_time(total_start_timer) + if not baseline_compare: + start_time(total_start_timer) for i in range(n_layers): if i == 0: nnz = batch_size @@ -64,7 +66,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, nnz = current_frontier[0, :].size(0) # Expand batches matrix - total_start_timer.record() + if baseline_compare: + total_start_timer.record() batches_expand_rows = torch.arange(mb_count * nnz, \ device=torch.device("cuda:0")) batches_expand_idxs = torch.stack( @@ -82,10 +85,17 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, timing_dict, "sage") - total_stop_timer.record() - torch.cuda.synchronize() - total_time = total_start_timer.elapsed_time(total_stop_timer) - print(f"total_time: {total_time}", flush=True) + + if baseline_compare: + total_stop_timer.record() + torch.cuda.synchronize() + total_time = total_start_timer.elapsed_time(total_stop_timer) + print(f"total_time: {total_time}", flush=True) + + batches_select = None + next_frontier_select = None + adj_matrices = None + break # for comparing with quiver # add explicit 0's to next_frontier next_frontier_nnz = next_frontier._values().nonzero().squeeze() @@ -123,6 +133,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier + print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) 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": # print(f"{k} times: {v}") diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index b1cf433..bd8a7f6 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -100,6 +100,8 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, # print(f"mata.size(0): {mata.size(0)}", flush=True) # mata_chunk_indices = mata_chunk_indices.int() # matb_recv_indices = matb_recv_indices.int() + # print(f"matb_recv_indices.row: {matb_recv_indices[0,:]}", flush=True) + # print(f"matb_recv_values: {matb_recv_values}", flush=True) # matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows # mata_chunk_indices[1,:], # A_cols # mata_chunk_values.float(), # A_vals @@ -110,6 +112,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, # chunk_col_size, # k # matb.size(1)) # m # print(f"after spgemm_gpu", flush=True) + # print(f"matc_outputs: {matc_outputs}", flush=True) # mata_chunk_indices = mata_chunk_indices.long() # matb_recv_indices = matb_recv_indices.long() # matc_chunk_rows = matc_outputs[0].long() @@ -353,47 +356,61 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_indices = matb_select_recv[:2, :].long() matb_recv_values = matb_select_recv[2, :].double() - start_time(start_timer) - # 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) - mata_chunk_indices = mata_chunk_indices.int() - matb_recv_indices = matb_recv_indices.int() - matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows - mata_chunk_indices[1,:], # A_cols - mata_chunk_values.float(), # A_vals - matb_recv_indices[0,:], # B_rows - matb_recv_indices[1,:], # B_cols - matb_recv_values.float(), # B_vals - mata.size(0), # n - chunk_col_size, # k - matb.size(1)) # m - mata_chunk_indices = mata_chunk_indices.long() - matb_recv_indices = matb_recv_indices.long() - 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, mata.size(0), device=torch.device("cuda:0")), - matc_chunk_counts) - matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) - matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: + start_time(start_timer) + 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) + # mata_chunk_indices = mata_chunk_indices.int() + # matb_recv_indices = matb_recv_indices.int() + # matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows + # mata_chunk_indices[1,:], # A_cols + # mata_chunk_values.float(), # A_vals + # matb_recv_indices[0,:], # B_rows + # matb_recv_indices[1,:], # B_cols + # matb_recv_values.float(), # B_vals + # mata.size(0), # n + # chunk_col_size, # k + # matb.size(1)) # m + # mata_chunk_indices = mata_chunk_indices.long() + # matb_recv_indices = matb_recv_indices.long() + # 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, mata.size(0), device=torch.device("cuda:0")), + # matc_chunk_counts) + # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + + 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, + torch.Size([matc.size(0), matc.size(1)])) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + + start_time(start_timer) + matc_chunk = matc_chunk.coalesce() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + + start_time(start_timer) + matc += matc_chunk + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + else: + if f"spgemm-local-spgemm-{name}" not in timing_dict: + timing_dict[f"spgemm-local-spgemm-{name}"] = [] - 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, - torch.Size([matc.size(0), matc.size(1)])) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + if f"spgemm-chunk-inst-{name}" not in timing_dict: + timing_dict[f"spgemm-chunk-inst-{name}"] = [] - start_time(start_timer) - matc_chunk = matc_chunk.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + if f"spgemm-chunk-add-{name}" not in timing_dict: + timing_dict[f"spgemm-chunk-add-{name}"] = [] - start_time(start_timer) - matc += matc_chunk - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + if f"spgemm-chunk-coalesce-{name}" not in timing_dict: + timing_dict[f"spgemm-chunk-coalesce-{name}"] = [] start_time(start_timer) matc = matc.coalesce() @@ -505,21 +522,27 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, n_darts_col = n_darts - (replication - 1) * n_darts_col n_darts_col = n_darts + next_frontier = torch.sparse_coo_tensor(indices=p._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) torch.cuda.nvtx.range_push("nvtx-pre-loop") - # next_frontier = torch.sparse_coo_tensor(indices=p._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])) + # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], bins=p.size(0)).int() - frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], - bins=p.size(0)).int() - # sampled_count = torch.cuda.IntTensor(p.size(0)).fill_(0) - sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) + frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0)).fill_(0) + frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], next_frontier._indices()[0,:].int()) + frontier_nnz_sizes = torch.div(frontier_nnz_sizes, + torch.arange(0, p.size(0), device=torch.device("cuda:0")), + rounding_mode="floor").int() + zero_count = (next_frontier._indices()[0,:] == 0).nonzero().size(0) + frontier_nnz_sizes[0] = zero_count torch.cuda.nvtx.range_pop() timing_dict["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 @@ -615,7 +638,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(select_start_timer) torch.cuda.nvtx.range_push("nvtx-dart-selection") - if rank_col == 0: + if rank_col == 0 and overflowed_minibatches: while overflowed_minibatches: start_time(select_iter_start_timer) torch.cuda.nvtx.range_push("nvtx-selection-iter") @@ -893,7 +916,7 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ # 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.size(1) * batch_size])) + 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) torch.cuda.nvtx.range_pop() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 498ad30..0a1f2f5 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -302,7 +302,7 @@ def main(args, batches=None): size = dist.get_world_size() print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) - adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) edge_count = adj_matrix.size(1) partitioning = Partitioning.ONE5D @@ -351,7 +351,7 @@ def main(args, batches=None): print("beginning constructing batches") batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - torch.manual_seed(0) + # torch.manual_seed(0) vertex_perm = torch.randperm(train_nid.size(0)) # Generate minibatch vertices for i in range(args.n_bulkmb): @@ -388,7 +388,7 @@ def main(args, batches=None): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - torch.manual_seed(0) + # torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) @@ -396,7 +396,7 @@ def main(args, batches=None): sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) if args.sample_method == "ladies": - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -407,7 +407,7 @@ def main(args, batches=None): print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices_bulk = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -415,7 +415,7 @@ def main(args, batches=None): args.replication, nnz_row_masks, sa_recv_buff, \ rank, size, row_groups, col_groups) elif args.sample_method == "sage": - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices = \ sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -426,7 +426,7 @@ def main(args, batches=None): print() torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices_bulk = \ sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -437,44 +437,54 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() - # adj_matrices[i][j] -- mb i layer j - rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) - adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] - for i in range(args.n_layers): - for j in range(rank_n_bulkmb): - row_select_min = j * args.batch_size - row_select_max = (j + 1) * args.batch_size - - sampled_indices = adj_matrices_bulk[i]._indices() - sampled_values = adj_matrices_bulk[i]._values() - - sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask] - - if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) - else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) - adj_matrices[j][i] = adj_matrix_sample - - # print(f"sample: {adj_matrices}") - # 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() - adj_matrix = row_normalize(adj_matrix) - adj_matrix = torch.pow(adj_matrix, 2) - # return here while testing sampling code - return current_frontier, next_frontier, adj_matrices, adj_matrix + if current_frontier is not None: + # adj_matrices[i][j] -- mb i layer j + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] + print(f"before reindexing", flush=True) + for i in range(args.n_layers): + for j in range(rank_n_bulkmb): + row_select_min = j * args.batch_size + row_select_max = (j + 1) * args.batch_size + + sampled_indices = adj_matrices_bulk[i]._indices() + sampled_values = adj_matrices_bulk[i]._values() + + sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask] + + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) + adj_matrices[j][i] = adj_matrix_sample + + print(f"after reindexing", flush=True) + # print(f"sample: {adj_matrices}") + # 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() + adj_matrix = row_normalize(adj_matrix) + adj_matrix = torch.pow(adj_matrix, 2) + # return here while testing sampling code + + print(f"current_frontier: {current_frontier}") + + print(f"next_frontier[20]: {next_frontier[(next_frontier[:,-1] == 20).nonzero().squeeze(), :]}") + print(f"next_frontier: {next_frontier}") + print(f"adj_matrices: {adj_matrices}") + return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + return # create GCN model torch.manual_seed(0) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index e18e4fc..b3e3c93 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -312,6 +312,7 @@ std::vector spgemm_gpu( 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, (n + 1) * sizeof(int32_t)); CHECK_CUSPARSE(cusparseXcoo2csr(handle, From 04ba929ec7d12e5113b7876c1087832f48042219 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 19 Oct 2022 17:51:46 -0700 Subject: [PATCH 104/178] Added coogeam extension to reduce reduction time --- .gitignore | 3 + cagnet/samplers/ladies.py | 7 +- cagnet/samplers/sage.py | 61 ++++++---- cagnet/samplers/utils.py | 100 +++++++++++++--- examples/cagnet_exps.sh | 29 +++++ examples/gcn_15d.py | 85 ++++++++++---- sparse-extension/sparse_coo_tensor.cu | 162 ++++++++++++++++++++++++++ 7 files changed, 379 insertions(+), 68 deletions(-) create mode 100755 examples/cagnet_exps.sh diff --git a/.gitignore b/.gitignore index 3cc0263..caa8558 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,8 @@ sparse-extension/sparse_coo_tensor_cpp.egg-info/ sparse_coo_tensor_cpp.egg-info/ .*.swp .*.swo +.*.swm +.*.swn .nfs* *.txt !requirements.txt @@ -31,3 +33,4 @@ tests/*.out *.qdrep *.out *.nsys-rep +*.ncu-rep diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 38bffbb..24e2148 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -11,6 +11,7 @@ shift_rowselect_gpu, shift_colselect_gpu, \ scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu + def start_time(timer): timer.record() @@ -18,9 +19,9 @@ def stop_time(start_timer, stop_timer): stop_timer.record() torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) - + def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ - replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups): + replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups, timing): total_start_timer = torch.cuda.Event(enable_timing=True) total_stop_timer = torch.cuda.Event(enable_timing=True) @@ -52,7 +53,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota nnz = current_frontier[0, :].size(0) p = gen_prob_dist(batches, adj_matrix, mb_count, node_count_total, replication, rank, size, \ - row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies") + row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies", timing) next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ row_groups, col_groups, timing_dict, "ladies") diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 9c8cbaf..f455115 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -14,14 +14,14 @@ timing = True -baseline_compare = False +baseline_compare = True def start_time(timer): - if timing == True: + if timing: timer.record() def stop_time(start_timer, stop_timer): - if timing == True: + if timing: stop_timer.record() torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) @@ -29,7 +29,14 @@ def stop_time(start_timer, stop_timer): return 0.0 def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ - replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups): + replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups, + timing_arg, baseline_arg): + + 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) @@ -58,7 +65,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrices = [None] * n_layers # adj_matrices[i] -- bulk mb mtx for layer j if not baseline_compare: - start_time(total_start_timer) + # start_time(total_start_timer) + total_start_timer.record() for i in range(n_layers): if i == 0: nnz = batch_size @@ -80,7 +88,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, - sa_masks, sa_recv_buff, timing_dict, "sage") + sa_masks, sa_recv_buff, timing_dict, "sage", + timing_arg) next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, @@ -133,24 +142,28 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier - print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) - 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": - # 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() + # print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) + total_stop_timer.record() + torch.cuda.synchronize() + total_time = total_start_timer.elapsed_time(total_stop_timer) + print(f"total_time: {total_time}", flush=True) + # 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": + # # 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) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index bd8a7f6..f269097 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1,7 +1,7 @@ -import math -import torch -import torch.distributed as dist -import torch_sparse +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, \ @@ -9,17 +9,17 @@ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ - sparse_coo_tensor_gpu, spgemm_gpu + sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu timing = True def start_time(timer): - if timing == True: + if timing: timer.record() def stop_time(start_timer, stop_timer): - if timing == True: + if timing: stop_timer.record() torch.cuda.synchronize() return start_timer.elapsed_time(stop_timer) @@ -200,6 +200,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_inner_timer = torch.cuda.Event(enable_timing=True) stop_inner_timer = torch.cuda.Event(enable_timing=True) + torch.cuda.nvtx.range_push("nvtx-matc-inst") 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() @@ -209,12 +210,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + torch.cuda.nvtx.range_pop() # matc-inst for i in range(stages): - start_time(start_timer) + # start_time(start_timer) q = (rank_col * (size // (replication ** 2)) + i) * replication + rank_col # Extract chunk of mata for local SpGEMM + torch.cuda.nvtx.range_push("nvtx-preproc-local") start_time(start_timer) am_partid = rank_col * (size // replication ** 2) + i chunk_col_start = am_partid * chunk_size @@ -226,14 +229,18 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + torch.cuda.nvtx.range_pop() # preproc-local # Determine number of nonzero columns in chunk of mata + torch.cuda.nvtx.range_push("nvtx-unique") start_time(start_timer) nnz_cols = torch.unique(mata_chunk_indices[1,:]) 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}") + torch.cuda.nvtx.range_pop() # unique # Gather nnz column counts on rank q + torch.cuda.nvtx.range_push("nvtx-gather-nnzcounts") start_time(start_timer) if rank == q: nnz_cols_count_list = [] @@ -244,16 +251,20 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + torch.cuda.nvtx.range_pop() # gather-nnzcounts # Rank q allocates recv buffer for nnz col ids + torch.cuda.nvtx.range_push("nvtx-alloc-nnzbuff") 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}") + torch.cuda.nvtx.range_pop() # alloc-nnzbuff # isend/irecv nnz col ids to rank q + torch.cuda.nvtx.range_push("nvtx-send-nnzcols") start_time(start_timer) if rank == q: recv_objs = [] @@ -266,9 +277,10 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups else: dist.isend(nnz_cols, dst=q) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") + torch.cuda.nvtx.range_pop() # send-nnzcols # Send selected rows count - torch.cuda.nvtx.range_push("nvtx-send-nnzrows") + torch.cuda.nvtx.range_push("nvtx-send-rowcounts") start_time(start_timer) if rank == q: # Add timings to timing_dict for src process for timing printouts @@ -286,6 +298,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups for j in range(size // replication): recv_rank = rank_col + j * replication + torch.cuda.nvtx.range_push("nvtx-send-ovhd") start_time(start_inner_timer) # start_time(start_inner_timer) @@ -318,6 +331,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") + torch.cuda.nvtx.range_pop() start_time(start_inner_timer) if recv_rank != q: @@ -351,12 +365,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") - torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() # send-rowcounts matb_recv_indices = matb_select_recv[:2, :].long() matb_recv_values = matb_select_recv[2, :].double() if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: + torch.cuda.nvtx.range_push("nvtx-local-spgemm") start_time(start_timer) matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ @@ -385,20 +400,27 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + torch.cuda.nvtx.range_pop() # local-spgemm + torch.cuda.nvtx.range_push("nvtx-chunk-inst") 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, torch.Size([matc.size(0), matc.size(1)])) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + torch.cuda.nvtx.range_pop() # chunk-inst + torch.cuda.nvtx.range_push("nvtx-chunk-coalesce") start_time(start_timer) matc_chunk = matc_chunk.coalesce() stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + torch.cuda.nvtx.range_pop() # chunk-coalesce + torch.cuda.nvtx.range_push("nvtx-chunk-add") start_time(start_timer) matc += matc_chunk stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + torch.cuda.nvtx.range_pop() # chunk-add else: if f"spgemm-local-spgemm-{name}" not in timing_dict: timing_dict[f"spgemm-local-spgemm-{name}"] = [] @@ -412,18 +434,23 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups if f"spgemm-chunk-coalesce-{name}" not in timing_dict: timing_dict[f"spgemm-chunk-coalesce-{name}"] = [] + torch.cuda.nvtx.range_push("nvtx-matc-coalesce") start_time(start_timer) matc = matc.coalesce() stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") + torch.cuda.nvtx.range_pop() # matc-coalesce # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) # Implement sparse allreduce w/ all_gather and padding + torch.cuda.nvtx.range_push("nvtx-reduce-nnz") 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}") + torch.cuda.nvtx.range_pop() # reduce-nnz + torch.cuda.nvtx.range_push("nvtx-padding") start_time(start_timer) matc_recv_indices = [] matc_recv_values = [] @@ -431,42 +458,78 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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))) + 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}") + 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}") + 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())) - matc_recv.append(sparse_coo_tensor_gpu(matc_recv_indices[i], matc_recv_values[i], matc.size())) + 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}") + 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) + # matc_recv = torch.stack(matc_recv) + # matc = torch.sparse.sum(matc_recv, dim=0) + + 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=torch.device("cuda:0")), + 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}") + 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}") + 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}") + 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, sa_recv_buff, - timing_dict, name): + timing_dict, name, timing_arg): + + global timing + timing = timing_arg start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -494,7 +557,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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])) - print(f"p.nnz: {p._nnz()}") + print(f"p.nnz: {p._nnz()}", flush=True) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) return p @@ -646,6 +709,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-psoverflow") selection_iter_count += 1 + print(f"selection_iter_count: {selection_iter_count}", flush=True) ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() # n_darts_select = total_overflow // replication 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/gcn_15d.py b/examples/gcn_15d.py index 0a1f2f5..c1aa30b 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -4,6 +4,7 @@ import os.path as osp import time import numpy as np + import torch import torch.distributed as dist import torch.nn as nn @@ -20,6 +21,9 @@ from cagnet.samplers import ladies_sampler, sage_sampler import cagnet.nn.functional as CAGF +import ogb +from ogb.nodeproppred import PygNodePropPredDataset + import socket class GCN(nn.Module): @@ -135,7 +139,9 @@ def row_normalize(mx): 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 = 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()], @@ -283,6 +289,29 @@ def main(args, batches=None): data = data.to(device) inputs.requires_grad = True data.y = data.y.to(device) + elif args.dataset.startswith("ogb"): + dataset = PygNodePropPredDataset(name=args.dataset, root="../../data") + + 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 = dataset[0] + # data = data.to(device) + # data.x.requires_grad = True + # inputs = data.x.to(device) + inputs = data.x + # 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 + split_idx = dataset.get_idx_split() + train_idx = split_idx['train'].to(device) + + print(f"adj_matrix.size(): {adj_matrix.size()}") + print(f"adj_matrix: {adj_matrix}") # Initialize distributed environment with SLURM if "SLURM_PROCID" in os.environ.keys(): @@ -318,11 +347,10 @@ def main(args, batches=None): features_loc, g_loc, ampbyp = one5d_partition(rank, size, inputs, adj_matrix, data, \ inputs, num_classes, args.replication, \ args.normalize) - del inputs print("end partitioning", flush=True) - features_loc = features_loc.to(device) + # 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) @@ -336,22 +364,25 @@ def main(args, batches=None): n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) - rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] - rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] - labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] - rank_train_nids = rank_train_mask.nonzero().squeeze() - rank_test_nids = rank_test_mask.nonzero().squeeze() - - train_nid = data.train_mask.nonzero().squeeze() - test_nid = data.test_mask.nonzero().squeeze() - + # rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] + # rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] + # labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] + # rank_train_nids = rank_train_mask.nonzero().squeeze() + # rank_test_nids = rank_test_mask.nonzero().squeeze() + + if not args.dataset.startswith("ogb"): + train_nid = data.train_mask.nonzero().squeeze() + test_nid = data.test_mask.nonzero().squeeze() + else: + train_nid = train_idx + torch.cuda.nvtx.range_pop() if batches is None: print("beginning constructing batches") batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - # torch.manual_seed(0) + torch.manual_seed(0) vertex_perm = torch.randperm(train_nid.size(0)) # Generate minibatch vertices for i in range(args.n_bulkmb): @@ -388,7 +419,7 @@ def main(args, batches=None): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}") # do it once before timing - # torch.manual_seed(0) + torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) @@ -401,8 +432,9 @@ def main(args, batches=None): ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + args.replication, nnz_row_masks, + sa_recv_buff, rank, size, row_groups, + col_groups, args.timing) print() torch.cuda.profiler.cudart().cudaProfilerStart() @@ -412,16 +444,18 @@ def main(args, batches=None): ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + args.replication, nnz_row_masks, + sa_recv_buff, rank, size, row_groups, + col_groups, args.timing) elif args.sample_method == "sage": - # torch.manual_seed(rank_col) + torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices = \ sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + args.replication, nnz_row_masks, + sa_recv_buff, rank, size, row_groups, + col_groups, args.timing, args.baseline) print() torch.cuda.profiler.cudart().cudaProfilerStart() @@ -431,8 +465,9 @@ def main(args, batches=None): sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, sa_recv_buff, \ - rank, size, row_groups, col_groups) + args.replication, nnz_row_masks, + sa_recv_buff, rank, size, row_groups, + col_groups, args.timing, args.baseline) torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() @@ -605,6 +640,10 @@ def main(args, batches=None): help='number of minibatches to sample 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('--timing', action="store_true", + help='whether to turn on timers') + parser.add_argument('--baseline', action="store_true", + help='whether to avoid col selection for baseline comparison') args = parser.parse_args() print(args) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index b3e3c93..172d38a 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -290,6 +290,167 @@ void spmm_gpu(const at::Tensor& A_rowindices, 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 + CHECK_ERROR("after 2nnz"); + 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 spgemm_gpu( const at::Tensor& A_rowindices, @@ -909,4 +1070,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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, "Row selection for sparsity-aware spgemm"); + m.def("coogeam_gpu", &coogeam_gpu, "csrgeam2 wrapper for cusparse"); } From 85139141e19c032477a1dbb77e5396ce926afa35 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 21 Nov 2022 07:56:02 -0800 Subject: [PATCH 105/178] New timer --- cagnet/samplers/ladies.py | 7 +- cagnet/samplers/sage.py | 41 ++++--- cagnet/samplers/utils.py | 161 ++++++++++++++------------ examples/gcn_15d.py | 16 +++ examples/job_sage | 20 ++++ examples/swp | 154 ++++++++++++++++++++++++ sparse-extension/sparse_coo_tensor.cu | 2 +- 7 files changed, 307 insertions(+), 94 deletions(-) create mode 100755 examples/job_sage create mode 100644 examples/swp diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 24e2148..f86aa47 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -15,10 +15,13 @@ def start_time(timer): timer.record() -def stop_time(start_timer, stop_timer): +def stop_time(start_timer, stop_timer, barrier=False): stop_timer.record() torch.cuda.synchronize() - return start_timer.elapsed_time(stop_timer) + 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, \ replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups, timing): diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index f455115..cf032e2 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -20,11 +20,14 @@ def start_time(timer): if timing: timer.record() -def stop_time(start_timer, stop_timer): +def stop_time(start_timer, stop_timer, barrier=False): if timing: stop_timer.record() torch.cuda.synchronize() - return start_timer.elapsed_time(stop_timer) + time_taken = start_timer.elapsed_time(stop_timer) + if barrier: + dist.barrier() + return time_taken else: return 0.0 @@ -147,23 +150,23 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - # 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": - # # 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 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": + # 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) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index f269097..0e77d95 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -18,15 +18,18 @@ def start_time(timer): if timing: timer.record() -def stop_time(start_timer, stop_timer): +def stop_time(start_timer, stop_timer, barrier=False): if timing: stop_timer.record() torch.cuda.synchronize() - return start_timer.elapsed_time(stop_timer) + time_taken = start_timer.elapsed_time(stop_timer) + if barrier: + dist.barrier() + return time_taken else: return 0.0 -def stop_time_add(start_timer, stop_timer, timing_dict, range_name): +def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=False): if timing_dict is not None: # misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ # not range_name.startswith("spgemm-bcast-data") and \ @@ -36,7 +39,7 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name): if False and misc_range: timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) else: - timing_dict[range_name].append(stop_time(start_timer, stop_timer)) + timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, name, timing_dict=None): @@ -209,7 +212,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-inst-{name}", barrier=True) torch.cuda.nvtx.range_pop() # matc-inst for i in range(stages): @@ -228,7 +231,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}", barrier=True) torch.cuda.nvtx.range_pop() # preproc-local # Determine number of nonzero columns in chunk of mata @@ -236,7 +239,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_time(start_timer) nnz_cols = torch.unique(mata_chunk_indices[1,:]) 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-unique-{name}", barrier=True) torch.cuda.nvtx.range_pop() # unique # Gather nnz column counts on rank q @@ -250,7 +253,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-gather-nnzcounts-{name}", barrier=True) torch.cuda.nvtx.range_pop() # gather-nnzcounts # Rank q allocates recv buffer for nnz col ids @@ -260,7 +263,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-alloc-nnzbuff-{name}", barrier=True) torch.cuda.nvtx.range_pop() # alloc-nnzbuff # isend/irecv nnz col ids to rank q @@ -276,36 +279,38 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_col_ids[j] = nnz_cols else: dist.isend(nnz_cols, dst=q) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-nnzcols # Send selected rows count torch.cuda.nvtx.range_push("nvtx-send-rowcounts") start_time(start_timer) if rank == q: - # Add timings to timing_dict for src process for timing printouts - if f"spgemm-recv-alloccount-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-alloccount-{name}"] = [] + # # Add timings to timing_dict for src process for timing printouts + # if f"spgemm-recv-alloccount-{name}" not in timing_dict: + # timing_dict[f"spgemm-recv-alloccount-{name}"] = [] - if f"spgemm-recv-rowcounts-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-rowcounts-{name}"] = [] + # if f"spgemm-recv-rowcounts-{name}" not in timing_dict: + # timing_dict[f"spgemm-recv-rowcounts-{name}"] = [] - if f"spgemm-recv-alloc-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-alloc-{name}"] = [] + # if f"spgemm-recv-alloc-{name}" not in timing_dict: + # timing_dict[f"spgemm-recv-alloc-{name}"] = [] - if f"spgemm-recv-rowdata-{name}" not in timing_dict: - timing_dict[f"spgemm-recv-rowdata-{name}"] = [] + # if f"spgemm-recv-rowdata-{name}" not in timing_dict: + # timing_dict[f"spgemm-recv-rowdata-{name}"] = [] for j in range(size // replication): recv_rank = rank_col + j * replication torch.cuda.nvtx.range_push("nvtx-send-ovhd") - start_time(start_inner_timer) + # start_time(start_inner_timer) # start_time(start_inner_timer) nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") # start_time(start_inner_timer) + if matb._indices().size(1) == 0: + print(f"nnz_col_count: {nnz_col_ids[j].size(0)} row_count: {matb._indices().size(1)}") rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ nnz_col_ids[j].size(0), matb._indices().size(1)) # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") @@ -330,49 +335,53 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") torch.cuda.nvtx.range_pop() - start_time(start_inner_timer) + # start_time(start_inner_timer) if recv_rank != q: selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) dist.isend(selected_rows_count, tag=0, dst=recv_rank) dist.isend(matb_send, tag=1, dst=recv_rank) else: matb_select_recv = matb_send.clone() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") else: - if f"spgemm-send-ovhd-{name}" not in timing_dict: - timing_dict[f"spgemm-send-ovhd-{name}"] = [] + # if f"spgemm-send-ovhd-{name}" not in timing_dict: + # timing_dict[f"spgemm-send-ovhd-{name}"] = [] - if f"spgemm-send-isend-{name}" not in timing_dict: - timing_dict[f"spgemm-send-isend-{name}"] = [] + # if f"spgemm-send-isend-{name}" not in timing_dict: + # timing_dict[f"spgemm-send-isend-{name}"] = [] - start_time(start_inner_timer) + # start_time(start_inner_timer) selected_rows_count_recv = torch.cuda.IntTensor(1) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") - start_time(start_inner_timer) + #start_time(start_inner_timer) dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowcounts-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowcounts-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") - start_time(start_inner_timer) + # start_time(start_inner_timer) dist.irecv(matb_select_recv, tag=1, src=q).wait() - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-rowcounts + torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") 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) + torch.cuda.nvtx.range_push("nvtx-local-spgemm") + start_time(start_timer) if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: - torch.cuda.nvtx.range_push("nvtx-local-spgemm") - start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-local-spgemm") + # start_time(start_timer) matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ matb_recv_values, mata.size(0), \ @@ -399,45 +408,48 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk_counts) # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") - torch.cuda.nvtx.range_pop() # local-spgemm + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + # torch.cuda.nvtx.range_pop() # local-spgemm - torch.cuda.nvtx.range_push("nvtx-chunk-inst") - start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-chunk-inst") + # 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, torch.Size([matc.size(0), matc.size(1)])) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") - torch.cuda.nvtx.range_pop() # chunk-inst + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") + # torch.cuda.nvtx.range_pop() # chunk-inst - torch.cuda.nvtx.range_push("nvtx-chunk-coalesce") - start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-chunk-coalesce") + # start_time(start_timer) matc_chunk = matc_chunk.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") - torch.cuda.nvtx.range_pop() # chunk-coalesce + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") + # torch.cuda.nvtx.range_pop() # chunk-coalesce - torch.cuda.nvtx.range_push("nvtx-chunk-add") - start_time(start_timer) + # torch.cuda.nvtx.range_push("nvtx-chunk-add") + # start_time(start_timer) matc += matc_chunk - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") - torch.cuda.nvtx.range_pop() # chunk-add + # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") + # torch.cuda.nvtx.range_pop() # chunk-add else: - if f"spgemm-local-spgemm-{name}" not in timing_dict: - timing_dict[f"spgemm-local-spgemm-{name}"] = [] + pass + # if f"spgemm-local-spgemm-{name}" not in timing_dict: + # timing_dict[f"spgemm-local-spgemm-{name}"] = [] - if f"spgemm-chunk-inst-{name}" not in timing_dict: - timing_dict[f"spgemm-chunk-inst-{name}"] = [] + # if f"spgemm-chunk-inst-{name}" not in timing_dict: + # timing_dict[f"spgemm-chunk-inst-{name}"] = [] - if f"spgemm-chunk-add-{name}" not in timing_dict: - timing_dict[f"spgemm-chunk-add-{name}"] = [] + # if f"spgemm-chunk-add-{name}" not in timing_dict: + # timing_dict[f"spgemm-chunk-add-{name}"] = [] - if f"spgemm-chunk-coalesce-{name}" not in timing_dict: - timing_dict[f"spgemm-chunk-coalesce-{name}"] = [] + # if f"spgemm-chunk-coalesce-{name}" not in timing_dict: + # timing_dict[f"spgemm-chunk-coalesce-{name}"] = [] + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) + torch.cuda.nvtx.range_pop() # local-spgemm torch.cuda.nvtx.range_push("nvtx-matc-coalesce") start_time(start_timer) matc = matc.coalesce() - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}", barrier=True) torch.cuda.nvtx.range_pop() # matc-coalesce # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) @@ -447,7 +459,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-nnz-{name}", barrier=True) torch.cuda.nvtx.range_pop() # reduce-nnz torch.cuda.nvtx.range_push("nvtx-padding") @@ -460,23 +472,29 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + 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}") + 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())) - # 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}") + # 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") @@ -504,13 +522,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + 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}") + 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") @@ -519,7 +537,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") + 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 @@ -541,7 +559,6 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, sa_recv_buff, timing_dict) - torch.cuda.nvtx.range_pop() timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) @@ -557,6 +574,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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])) + # print(f"p: {p}") print(f"p.nnz: {p._nnz()}", flush=True) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) @@ -709,7 +727,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-select-psoverflow") selection_iter_count += 1 - print(f"selection_iter_count: {selection_iter_count}", flush=True) ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() # n_darts_select = total_overflow // replication diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index c1aa30b..8c30dc1 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -289,6 +289,22 @@ def main(args, batches=None): 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.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.startswith("ogb"): dataset = PygNodePropPredDataset(name=args.dataset, root="../../data") diff --git a/examples/job_sage b/examples/job_sage new file mode 100755 index 0000000..32821e0 --- /dev/null +++ b/examples/job_sage @@ -0,0 +1,20 @@ +#!/bin/bash +#SBATCH --qos=regular +#SBATCH --time=00:30:00 +#SBATCH --nodes=9 +#SBATCH --gpus=36 +#SBATCH --ntasks-per-node=4 +#SBATCH --constraint=gpu +#SBATCH --account=m1982_g + +module load pytorch +export PYTHONPATH="/global/homes/a/alokt/CAGNET:$PYTHONPATH" +export PYTHONPATH="/global/homes/a/alokt/CAGNET/sparse-extension/build/lib/python3.9/site-packages/sparse_coo_tensor_cpp-0.0.0-py3.9-linux-x86_64.egg:$PYTHONPATH" + +hostname=$(scontrol show hostnames | head -n 1) +echo $hostname +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_coogeam.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_coogeam.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_coogeam.out diff --git a/examples/swp b/examples/swp new file mode 100644 index 0000000..72da6f5 --- /dev/null +++ b/examples/swp @@ -0,0 +1,154 @@ +cagnet_exps.sh +gcn_15d.py +gcn_1d.py +gcn_2d.py +gcn_3d.py +__init__.py +nsysc_amazon_bs512sn10_mb256_sage_p4c2.ncu-rep +nsys_reddit_bs512sn10_mb64_sage_baseline.nsys-rep +report1.nsys-rep +report2.nsys-rep +report3.nsys-rep +report4.nsys-rep +test_spgemm.py +tmp_amazon_sage_bs128sn10_mb64_p1.out +tmp_amazon_sage_bs128sn10_mb64_p4c2.out +tmp_amazon_sage_bs256sn25_mb256_p1.out +tmp_amazon_sage_bs256sn25_mb256_p4c2.out +tmp_amazon_sage_bs256sn25_mb64_p1.out +tmp_amazon_sage_bs512sn10_mb128_p1.out +tmp_amazon_sage_bs512sn10_mb128_p2.out +tmp_amazon_sage_bs512sn10_mb128_p4c2.out +tmp_amazon_sage_bs512sn10_mb256_p1.out +tmp_amazon_sage_bs512sn10_mb256_p4c1_coogeam.out +tmp_amazon_sage_bs512sn10_mb256_p4c2_coalesce.out +tmp_amazon_sage_bs512sn10_mb256_p4c2_coogeam.out +tmp_amazon_sage_bs512sn10_mb256_p4c2_nocoalesce.out +tmp_amazon_sage_bs512sn10_mb256_p4c2.out +tmp_amazon_sage_bs512sn10_mb256_p4c2_sum.out +tmp_amazon_sage_bs512sn10_mb256_p4.out +tmp_amazon_sage_bs512sn10_mb512_p1.out +tmp_amazon_sage_bs512sn10_mb512_p4c2.out +tmp_amazon_sage_bs512sn10_mb64_p1.out +tmp_amazon_sage_bs512sn10_mb64_p4c2.out +tmp_cora_ladies_bs16sn16_mb1.out +tmp_cora_ladies_bs16sn16_mb2_p1.out +tmp_cora_ladies_bs16sn16_mb2_p2.out +tmp_cora_ladies_bs16sn4_mb1.out +tmp_cora_p1_2.out +tmp_cora_p1.out +tmp_cora_p2.out +tmp_cora_p4c1.out +tmp_cora_p4c2.out +tmp_cora_p4.out +tmp_cora_sage_bs16sn1_mb1.out +tmp_cora_sage_bs16sn2_mb1_2.out +tmp_cora_sage_bs16sn2_mb1.out +tmp_cora_sage_bs16sn2_mb2.out +tmp_cora_sage_bs16sn4_mb1.out +tmp_cora_sage_bs16sn4_mb1_p1_2.out +tmp_cora_sage_bs16sn4_mb1_p1.out +tmp_cora_sage_bs16sn4_mb2.out +tmp_cora_sage_bs16sn4_mb2_p1.out +tmp_cora_sage_bs16sn4_mb2_p2.out +tmp_cora_sage_bs16sn4_mb4_p1.out +tmp_cora_sage_bs16sn4_mb4_p2.out +tmp_cora_sage_bs16sn4_mb4_p4c2_coalesce.out +tmp_cora_sage_bs16sn4_mb4_p4c2_coogeam.out +tmp_cora_sage_bs16sn4_mb4_p4c2_noreduce.out +tmp_cora_sage_bs16sn4_mb4_p4c2.out +tmp_cora_sage_bs16sn4_mb4_p4c2_sum.out +tmp_cora_sage_bs16sn4_mb4_p4.out +tmp_cora_sage_bs16sn4_mb8_p1.out +tmp_cora_sage_bs16sn4_mb8_p2.out +tmp_cora_sage_bs16sn4_mb8_p4c2.out +tmp_cora_sage_bs16sn4_mb8_p4.out +tmp_p2c1_reddit.out +tmp_p2c1_reddit_saspgemm.out +tmp_p4c1_amazon_saspgemm.out +tmp_p4c1_amazon_saspgemm_scatteradd.out +tmp_p4c1_amazon_saspgemm_sepnodes.out +tmp_p4c1_amazon_spgemm.out +tmp_p4c1_reddit.out +tmp_p4c1_reddit_saspgemm_nnnzrowmasks.out +tmp_p4c1_reddit_saspgemm.out +tmp_p4c1_reddit_saspgemm_sarowsel.out +tmp_p4c1_reddit_saspgemm_scatteradd.out +tmp_p4c1_reddit_saspgemm_sepnodes.out +tmp_p4c1_reddit_spgemm.out +tmp_p4c1_reddit_spgemm_sepnodes.out +tmp_p4c2_amazon_saspgemm_scatteradd.out +tmp_p4c2_reddit.out +tmp_p4c2_reddit_saspgemm.out +tmp_p4c2_reddit_saspgemm_scatteradd.out +tmp_p8c1_amazon_saspgemm_scatteradd.out +tmp_p8c2_amazon_saspgemm_scatteradd.out +tmp_papers_ladies_bs256sn256_mb8_p8c2.out +tmp_papers_ladies_bs256sn256_mb8_p8.out +tmp_papers_sage_bs128sn25_mb8_p8.out +tmp_papers_sage_bs256sn256_mb8_p8.out +tmp_papers_sage_bs256sn25_mb16_p16.out +tmp_papers_sage_bs256sn25_mb8_p8.out +tmp_papers_sage_bs64sn25_mb8_p8.out +tmp_products_ladies_bs256sn256_mb8_p8c2.out +tmp_products_sage_bs256sn256_mb8_p8c2.out +tmp_products_sage_bs256sn25_mb128_p1.out +tmp_products_sage_bs256sn25_mb128_p4c2.out +tmp_products_sage_bs256sn25_mb128_p4.out +tmp_products_sage_bs256sn25_mb128_p8c2.out +tmp_products_sage_bs256sn25_mb128_p8.out +tmp_products_sage_bs256sn25_mb16_p8c2.out +tmp_products_sage_bs256sn25_mb1_p1.out +tmp_products_sage_bs256sn25_mb32_p8.out +tmp_products_sage_bs256sn25_mb64_p8c2.out +tmp_products_sage_bs256sn25_mb64_p8.out +tmp_products_sage_bs256sn25_mb8_p8c2.out +tmp_products_sage_bs256sn25_mb8_p8.out +tmp_reddit_ladies_bs256sn256_mb128_p4.out +tmp_reddit_ladies_bs256sn256_mb1_p1.out +tmp_reddit_ladies_bs256sn256_mb1_p4.out +tmp_reddit_sage_bs128sn10_mb64_baseline_p1.out +tmp_reddit_sage_bs128sn10_mb64_p1.out +tmp_reddit_sage_bs128sn10_mb64_timing_baseline_p1.out +tmp_reddit_sage_bs128sn10_mb64_timing_p1.out +tmp_reddit_sage_bs256sn10_mb256_p4.out +tmp_reddit_sage_bs256sn25_mb128_p1.out +tmp_reddit_sage_bs256sn25_mb128_p2.out +tmp_reddit_sage_bs256sn25_mb128_p4c2.out +tmp_reddit_sage_bs256sn25_mb128_p4.out +tmp_reddit_sage_bs256sn25_mb128_p8c2.out +tmp_reddit_sage_bs256sn25_mb128_p8.out +tmp_reddit_sage_bs256sn25_mb128_times_p4.out +tmp_reddit_sage_bs256sn25_mb16_p8c2.out +tmp_reddit_sage_bs256sn25_mb16_p8.out +tmp_reddit_sage_bs256sn25_mb1_p1.out +tmp_reddit_sage_bs256sn25_mb256_p1.out +tmp_reddit_sage_bs256sn25_mb256_p4c2.out +tmp_reddit_sage_bs256sn25_mb256_p4.out +tmp_reddit_sage_bs256sn25_mb4_p4.out +tmp_reddit_sage_bs256sn25_mb64_p1.out +tmp_reddit_sage_bs256sn25_mb64_p2.out +tmp_reddit_sage_bs256sn25_mb64_p4.out +tmp_reddit_sage_bs256sn25_mb8_p8c2.out +tmp_reddit_sage_bs256sn4_mb128_p4.out +tmp_reddit_sage_bs256sn4_mb1_p2.out +tmp_reddit_sage_bs256sn4_mb256_p4.out +tmp_reddit_sage_bs256sn4_mb2_p2.out +tmp_reddit_sage_bs256sn4_mb512_p1.out +tmp_reddit_sage_bs256sn4_mb512_p2.out +tmp_reddit_sage_bs256sn4_mb512_p4.out +tmp_reddit_sage_bs512sn10_mb256_p4c1_coogeam.out +tmp_reddit_sage_bs512sn10_mb256_p4c2_coogeam.out +tmp_reddit_sage_bs512sn10_mb256_p4c2_sum.out +tmp_reddit_sage_bs512sn10_mb64_baseline_p1.out +tmp_reddit_sage_bs512sn10_mb64_p1.out +tmp_reddit_sage_bs512sn10_mb64_timing_baseline_p1.out +tmp_reddit_sage_bs512sn10_mb64_timing_p1.out + +cagnet_outputs: +ogbn-products +reddit + +__pycache__: +gcn_15d.cpython-39.pyc +__init__.cpython-39.pyc diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 172d38a..ee57ca0 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1041,7 +1041,7 @@ void rowselect_coo_gpu(const at::Tensor& nnz_cols, const at::Tensor& rows, const int BLOCK_COUNT = std::ceil(row_count / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - if (nnz_col_count == 0) { + if (nnz_col_count == 0 || row_count == 0) { return; } From e8280961af28a3269596ec5002ab5859fb52b5ea Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 21 Nov 2022 13:36:27 -0800 Subject: [PATCH 106/178] Cleanup --- cagnet/samplers/sage.py | 1 - cagnet/samplers/utils.py | 138 +++-------------------------------- examples/swp | 154 --------------------------------------- 3 files changed, 11 insertions(+), 282 deletions(-) delete mode 100644 examples/swp diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index cf032e2..738f44f 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -152,7 +152,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, print(f"total_time: {total_time}", flush=True) 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": - # print(f"{k} times: {v}") v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) v_tens_recv = [] for i in range(size): diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 0e77d95..0a94d27 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -31,15 +31,7 @@ def stop_time(start_timer, stop_timer, barrier=False): def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=False): if timing_dict is not None: - # misc_range = range_name.startswith("spgemm-") and not range_name.startswith("spgemm-bcast-nnz") and \ - # not range_name.startswith("spgemm-bcast-data") and \ - # not range_name.startswith("spgemm-local-spgemm") and \ - # not range_name.startswith("spgemm-reduce-nnz") and \ - # not range_name.startswith("spgemm-reduce") - if False and misc_range: - timing_dict["spgemm-misc"].append(stop_time(start_timer, stop_timer)) - else: - timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) + timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, name, timing_dict=None): @@ -267,7 +259,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_pop() # alloc-nnzbuff # isend/irecv nnz col ids to rank q - torch.cuda.nvtx.range_push("nvtx-send-nnzcols") + torch.cuda.nvtx.range_push("nvtx-send-colids") start_time(start_timer) if rank == q: recv_objs = [] @@ -279,170 +271,63 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_col_ids[j] = nnz_cols else: dist.isend(nnz_cols, dst=q) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-nnzcols-{name}", barrier=True) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-colids-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-nnzcols # Send selected rows count - torch.cuda.nvtx.range_push("nvtx-send-rowcounts") + torch.cuda.nvtx.range_push("nvtx-send-rowdata") start_time(start_timer) if rank == q: - # # Add timings to timing_dict for src process for timing printouts - # if f"spgemm-recv-alloccount-{name}" not in timing_dict: - # timing_dict[f"spgemm-recv-alloccount-{name}"] = [] - - # if f"spgemm-recv-rowcounts-{name}" not in timing_dict: - # timing_dict[f"spgemm-recv-rowcounts-{name}"] = [] - - # if f"spgemm-recv-alloc-{name}" not in timing_dict: - # timing_dict[f"spgemm-recv-alloc-{name}"] = [] - - # if f"spgemm-recv-rowdata-{name}" not in timing_dict: - # timing_dict[f"spgemm-recv-rowdata-{name}"] = [] - for j in range(size // replication): recv_rank = rank_col + j * replication - torch.cuda.nvtx.range_push("nvtx-send-ovhd") - # start_time(start_inner_timer) - - # start_time(start_inner_timer) nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselectalc-{name}") - - # start_time(start_inner_timer) - if matb._indices().size(1) == 0: - print(f"nnz_col_count: {nnz_col_ids[j].size(0)} row_count: {matb._indices().size(1)}") + rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ nnz_col_ids[j].size(0), matb._indices().size(1)) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowselect-{name}") - # start_time(start_inner_timer) matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskindalc-{name}") - - # start_time(start_inner_timer) matb_send_indices = matb._indices()[:, nnz_row_mask] - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskind-{name}") - - # start_time(start_inner_timer) matb_send_indices = matb_send_indices.float() - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-indfloat-{name}") - # start_time(start_inner_timer) matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-maskval-{name}") - - # start_time(start_inner_timer) matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-cat-{name}") - - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-ovhd-{name}") - torch.cuda.nvtx.range_pop() - - # start_time(start_inner_timer) + if recv_rank != q: selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) dist.isend(selected_rows_count, tag=0, dst=recv_rank) dist.isend(matb_send, tag=1, dst=recv_rank) else: matb_select_recv = matb_send.clone() - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-isend-{name}") else: - # if f"spgemm-send-ovhd-{name}" not in timing_dict: - # timing_dict[f"spgemm-send-ovhd-{name}"] = [] - - # if f"spgemm-send-isend-{name}" not in timing_dict: - # timing_dict[f"spgemm-send-isend-{name}"] = [] - - # start_time(start_inner_timer) selected_rows_count_recv = torch.cuda.IntTensor(1) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloccount-{name}") - - #start_time(start_inner_timer) dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowcounts-{name}") - # start_time(start_inner_timer) matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-alloc-{name}") - - # start_time(start_inner_timer) dist.irecv(matb_select_recv, tag=1, src=q).wait() - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-recv-rowdata-{name}") - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowcounts-{name}", barrier=True) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowdata-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-rowcounts torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") + # 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, :].double() + matb_recv_values = matb_select_recv[2, :] stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-castrecv-{name}", barrier=True) torch.cuda.nvtx.range_push("nvtx-local-spgemm") start_time(start_timer) if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: - # torch.cuda.nvtx.range_push("nvtx-local-spgemm") - # start_time(start_timer) 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) - # mata_chunk_indices = mata_chunk_indices.int() - # matb_recv_indices = matb_recv_indices.int() - # matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows - # mata_chunk_indices[1,:], # A_cols - # mata_chunk_values.float(), # A_vals - # matb_recv_indices[0,:], # B_rows - # matb_recv_indices[1,:], # B_cols - # matb_recv_values.float(), # B_vals - # mata.size(0), # n - # chunk_col_size, # k - # matb.size(1)) # m - # mata_chunk_indices = mata_chunk_indices.long() - # matb_recv_indices = matb_recv_indices.long() - # 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, mata.size(0), device=torch.device("cuda:0")), - # matc_chunk_counts) - # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) - # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) - # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") - # torch.cuda.nvtx.range_pop() # local-spgemm - - # torch.cuda.nvtx.range_push("nvtx-chunk-inst") - # 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, torch.Size([matc.size(0), matc.size(1)])) - # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-inst-{name}") - # torch.cuda.nvtx.range_pop() # chunk-inst - - # torch.cuda.nvtx.range_push("nvtx-chunk-coalesce") - # start_time(start_timer) - matc_chunk = matc_chunk.coalesce() - # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-coalesce-{name}") - # torch.cuda.nvtx.range_pop() # chunk-coalesce - - # torch.cuda.nvtx.range_push("nvtx-chunk-add") - # start_time(start_timer) + # matc_chunk = matc_chunk.coalesce() matc += matc_chunk - # stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-chunk-add-{name}") - # torch.cuda.nvtx.range_pop() # chunk-add else: pass - # if f"spgemm-local-spgemm-{name}" not in timing_dict: - # timing_dict[f"spgemm-local-spgemm-{name}"] = [] - - # if f"spgemm-chunk-inst-{name}" not in timing_dict: - # timing_dict[f"spgemm-chunk-inst-{name}"] = [] - - # if f"spgemm-chunk-add-{name}" not in timing_dict: - # timing_dict[f"spgemm-chunk-add-{name}"] = [] - - # if f"spgemm-chunk-coalesce-{name}" not in timing_dict: - # timing_dict[f"spgemm-chunk-coalesce-{name}"] = [] stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) torch.cuda.nvtx.range_pop() # local-spgemm @@ -452,7 +337,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-matc-coalesce-{name}", barrier=True) torch.cuda.nvtx.range_pop() # matc-coalesce - # dist.all_reduce(matc, op=dist.reduce_op.SUM, group=row_groups[rank_c]) # Implement sparse allreduce w/ all_gather and padding torch.cuda.nvtx.range_push("nvtx-reduce-nnz") start_time(start_timer) diff --git a/examples/swp b/examples/swp deleted file mode 100644 index 72da6f5..0000000 --- a/examples/swp +++ /dev/null @@ -1,154 +0,0 @@ -cagnet_exps.sh -gcn_15d.py -gcn_1d.py -gcn_2d.py -gcn_3d.py -__init__.py -nsysc_amazon_bs512sn10_mb256_sage_p4c2.ncu-rep -nsys_reddit_bs512sn10_mb64_sage_baseline.nsys-rep -report1.nsys-rep -report2.nsys-rep -report3.nsys-rep -report4.nsys-rep -test_spgemm.py -tmp_amazon_sage_bs128sn10_mb64_p1.out -tmp_amazon_sage_bs128sn10_mb64_p4c2.out -tmp_amazon_sage_bs256sn25_mb256_p1.out -tmp_amazon_sage_bs256sn25_mb256_p4c2.out -tmp_amazon_sage_bs256sn25_mb64_p1.out -tmp_amazon_sage_bs512sn10_mb128_p1.out -tmp_amazon_sage_bs512sn10_mb128_p2.out -tmp_amazon_sage_bs512sn10_mb128_p4c2.out -tmp_amazon_sage_bs512sn10_mb256_p1.out -tmp_amazon_sage_bs512sn10_mb256_p4c1_coogeam.out -tmp_amazon_sage_bs512sn10_mb256_p4c2_coalesce.out -tmp_amazon_sage_bs512sn10_mb256_p4c2_coogeam.out -tmp_amazon_sage_bs512sn10_mb256_p4c2_nocoalesce.out -tmp_amazon_sage_bs512sn10_mb256_p4c2.out -tmp_amazon_sage_bs512sn10_mb256_p4c2_sum.out -tmp_amazon_sage_bs512sn10_mb256_p4.out -tmp_amazon_sage_bs512sn10_mb512_p1.out -tmp_amazon_sage_bs512sn10_mb512_p4c2.out -tmp_amazon_sage_bs512sn10_mb64_p1.out -tmp_amazon_sage_bs512sn10_mb64_p4c2.out -tmp_cora_ladies_bs16sn16_mb1.out -tmp_cora_ladies_bs16sn16_mb2_p1.out -tmp_cora_ladies_bs16sn16_mb2_p2.out -tmp_cora_ladies_bs16sn4_mb1.out -tmp_cora_p1_2.out -tmp_cora_p1.out -tmp_cora_p2.out -tmp_cora_p4c1.out -tmp_cora_p4c2.out -tmp_cora_p4.out -tmp_cora_sage_bs16sn1_mb1.out -tmp_cora_sage_bs16sn2_mb1_2.out -tmp_cora_sage_bs16sn2_mb1.out -tmp_cora_sage_bs16sn2_mb2.out -tmp_cora_sage_bs16sn4_mb1.out -tmp_cora_sage_bs16sn4_mb1_p1_2.out -tmp_cora_sage_bs16sn4_mb1_p1.out -tmp_cora_sage_bs16sn4_mb2.out -tmp_cora_sage_bs16sn4_mb2_p1.out -tmp_cora_sage_bs16sn4_mb2_p2.out -tmp_cora_sage_bs16sn4_mb4_p1.out -tmp_cora_sage_bs16sn4_mb4_p2.out -tmp_cora_sage_bs16sn4_mb4_p4c2_coalesce.out -tmp_cora_sage_bs16sn4_mb4_p4c2_coogeam.out -tmp_cora_sage_bs16sn4_mb4_p4c2_noreduce.out -tmp_cora_sage_bs16sn4_mb4_p4c2.out -tmp_cora_sage_bs16sn4_mb4_p4c2_sum.out -tmp_cora_sage_bs16sn4_mb4_p4.out -tmp_cora_sage_bs16sn4_mb8_p1.out -tmp_cora_sage_bs16sn4_mb8_p2.out -tmp_cora_sage_bs16sn4_mb8_p4c2.out -tmp_cora_sage_bs16sn4_mb8_p4.out -tmp_p2c1_reddit.out -tmp_p2c1_reddit_saspgemm.out -tmp_p4c1_amazon_saspgemm.out -tmp_p4c1_amazon_saspgemm_scatteradd.out -tmp_p4c1_amazon_saspgemm_sepnodes.out -tmp_p4c1_amazon_spgemm.out -tmp_p4c1_reddit.out -tmp_p4c1_reddit_saspgemm_nnnzrowmasks.out -tmp_p4c1_reddit_saspgemm.out -tmp_p4c1_reddit_saspgemm_sarowsel.out -tmp_p4c1_reddit_saspgemm_scatteradd.out -tmp_p4c1_reddit_saspgemm_sepnodes.out -tmp_p4c1_reddit_spgemm.out -tmp_p4c1_reddit_spgemm_sepnodes.out -tmp_p4c2_amazon_saspgemm_scatteradd.out -tmp_p4c2_reddit.out -tmp_p4c2_reddit_saspgemm.out -tmp_p4c2_reddit_saspgemm_scatteradd.out -tmp_p8c1_amazon_saspgemm_scatteradd.out -tmp_p8c2_amazon_saspgemm_scatteradd.out -tmp_papers_ladies_bs256sn256_mb8_p8c2.out -tmp_papers_ladies_bs256sn256_mb8_p8.out -tmp_papers_sage_bs128sn25_mb8_p8.out -tmp_papers_sage_bs256sn256_mb8_p8.out -tmp_papers_sage_bs256sn25_mb16_p16.out -tmp_papers_sage_bs256sn25_mb8_p8.out -tmp_papers_sage_bs64sn25_mb8_p8.out -tmp_products_ladies_bs256sn256_mb8_p8c2.out -tmp_products_sage_bs256sn256_mb8_p8c2.out -tmp_products_sage_bs256sn25_mb128_p1.out -tmp_products_sage_bs256sn25_mb128_p4c2.out -tmp_products_sage_bs256sn25_mb128_p4.out -tmp_products_sage_bs256sn25_mb128_p8c2.out -tmp_products_sage_bs256sn25_mb128_p8.out -tmp_products_sage_bs256sn25_mb16_p8c2.out -tmp_products_sage_bs256sn25_mb1_p1.out -tmp_products_sage_bs256sn25_mb32_p8.out -tmp_products_sage_bs256sn25_mb64_p8c2.out -tmp_products_sage_bs256sn25_mb64_p8.out -tmp_products_sage_bs256sn25_mb8_p8c2.out -tmp_products_sage_bs256sn25_mb8_p8.out -tmp_reddit_ladies_bs256sn256_mb128_p4.out -tmp_reddit_ladies_bs256sn256_mb1_p1.out -tmp_reddit_ladies_bs256sn256_mb1_p4.out -tmp_reddit_sage_bs128sn10_mb64_baseline_p1.out -tmp_reddit_sage_bs128sn10_mb64_p1.out -tmp_reddit_sage_bs128sn10_mb64_timing_baseline_p1.out -tmp_reddit_sage_bs128sn10_mb64_timing_p1.out -tmp_reddit_sage_bs256sn10_mb256_p4.out -tmp_reddit_sage_bs256sn25_mb128_p1.out -tmp_reddit_sage_bs256sn25_mb128_p2.out -tmp_reddit_sage_bs256sn25_mb128_p4c2.out -tmp_reddit_sage_bs256sn25_mb128_p4.out -tmp_reddit_sage_bs256sn25_mb128_p8c2.out -tmp_reddit_sage_bs256sn25_mb128_p8.out -tmp_reddit_sage_bs256sn25_mb128_times_p4.out -tmp_reddit_sage_bs256sn25_mb16_p8c2.out -tmp_reddit_sage_bs256sn25_mb16_p8.out -tmp_reddit_sage_bs256sn25_mb1_p1.out -tmp_reddit_sage_bs256sn25_mb256_p1.out -tmp_reddit_sage_bs256sn25_mb256_p4c2.out -tmp_reddit_sage_bs256sn25_mb256_p4.out -tmp_reddit_sage_bs256sn25_mb4_p4.out -tmp_reddit_sage_bs256sn25_mb64_p1.out -tmp_reddit_sage_bs256sn25_mb64_p2.out -tmp_reddit_sage_bs256sn25_mb64_p4.out -tmp_reddit_sage_bs256sn25_mb8_p8c2.out -tmp_reddit_sage_bs256sn4_mb128_p4.out -tmp_reddit_sage_bs256sn4_mb1_p2.out -tmp_reddit_sage_bs256sn4_mb256_p4.out -tmp_reddit_sage_bs256sn4_mb2_p2.out -tmp_reddit_sage_bs256sn4_mb512_p1.out -tmp_reddit_sage_bs256sn4_mb512_p2.out -tmp_reddit_sage_bs256sn4_mb512_p4.out -tmp_reddit_sage_bs512sn10_mb256_p4c1_coogeam.out -tmp_reddit_sage_bs512sn10_mb256_p4c2_coogeam.out -tmp_reddit_sage_bs512sn10_mb256_p4c2_sum.out -tmp_reddit_sage_bs512sn10_mb64_baseline_p1.out -tmp_reddit_sage_bs512sn10_mb64_p1.out -tmp_reddit_sage_bs512sn10_mb64_timing_baseline_p1.out -tmp_reddit_sage_bs512sn10_mb64_timing_p1.out - -cagnet_outputs: -ogbn-products -reddit - -__pycache__: -gcn_15d.cpython-39.pyc -__init__.cpython-39.pyc From 03cbdbf30068d9359b80298ad5a1ac10d498ce26 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 5 Dec 2022 13:58:49 -0800 Subject: [PATCH 107/178] New timers --- cagnet/samplers/utils.py | 71 ++++++++++++++++++++++++++++++++-------- examples/gcn_15d.py | 2 +- 2 files changed, 59 insertions(+), 14 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 0a94d27..ed69891 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -33,7 +33,6 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=Fals if timing_dict is not None: timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) - def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, name, timing_dict=None): start_timer = torch.cuda.Event(enable_timing=True) @@ -233,6 +232,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) torch.cuda.nvtx.range_pop() # unique + # timing_dict[f"spgemm-nnzcount-{name}"].append(nnz_cols.size(0)) # Gather nnz column counts on rank q torch.cuda.nvtx.range_push("nvtx-gather-nnzcounts") @@ -277,42 +277,84 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # Send selected rows count torch.cuda.nvtx.range_push("nvtx-send-rowdata") start_time(start_timer) + nnz_count = 0 if rank == q: for j in range(size // replication): + # start_time(start_inner_timer) recv_rank = rank_col + j * replication nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] - rowselect_coo_gpu(nnz_col_ids[j].long(), matb._indices()[0,:], nnz_row_mask, \ + rowselect_coo_gpu(nnz_col_ids[j], matb._indices()[0,:], nnz_row_mask, \ nnz_col_ids[j].size(0), matb._indices().size(1)) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowselect-{name}") + + # matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) + # torch.cuda.synchronize() - matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) + # start_time(start_inner_timer) matb_send_indices = matb._indices()[:, nnz_row_mask] - matb_send_indices = matb_send_indices.float() + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowind-{name}") + + # start_time(start_inner_timer) + matb_send_values = matb._values()[nnz_row_mask] + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowval-{name}") - matb_send_values = matb._values()[nnz_row_mask].unsqueeze(dim=0) - matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) + # start_time(start_inner_timer) + # # matb_send_values = matb_send_values.unsqueeze(dim=0) + # # matb_send_indices = matb_send_indices.float() + # # matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowmisc-{name}") if recv_rank != q: - selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) + # selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_indices.size(1)) + nnz_count += 3 * matb_send_indices.size(1) dist.isend(selected_rows_count, tag=0, dst=recv_rank) - dist.isend(matb_send, tag=1, dst=recv_rank) + # dist.isend(matb_send, tag=1, dst=recv_rank) + dist.isend(matb_send_indices, tag=1, dst=recv_rank) + dist.isend(matb_send_values, tag=2, dst=recv_rank) else: - matb_select_recv = matb_send.clone() + # matb_select_recv = matb_send.clone() + matb_recv_indices = matb_send_indices.clone() + matb_recv_values = matb_send_values.clone() else: + # if f"spgemm-rowselect-{name}" not in timing_dict: + # timing_dict[f"spgemm-rowselect-{name}"] = [] + + # if f"spgemm-rowind-{name}" not in timing_dict: + # timing_dict[f"spgemm-rowind-{name}"] = [] + + # if f"spgemm-rowval-{name}" not in timing_dict: + # timing_dict[f"spgemm-rowval-{name}"] = [] + selected_rows_count_recv = torch.cuda.IntTensor(1) dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() - matb_select_recv = matb_recv_buff[:(3 *selected_rows_count_recv.item())].view(3, -1) - dist.irecv(matb_select_recv, tag=1, src=q).wait() + # 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()) + ind_obj = dist.irecv(matb_recv_indices, tag=1, src=q) + + matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) + val_obj = dist.irecv(matb_recv_values, tag=2, src=q) + + ind_obj.wait() + val_obj.wait() + + nnz_count += 3 * matb_recv_indices.size(1) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowdata-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-rowcounts + # timing_dict[f"spgemm-send-rownnz-{name}"].append(nnz_count) torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") + start_time(start_timer) # print(f"matb_select_recv.dtype: {matb_select_recv.dtype}") - matb_recv_indices = matb_select_recv[:2, :].long() + # 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() - matb_recv_values = matb_select_recv[2, :] stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-castrecv-{name}", barrier=True) torch.cuda.nvtx.range_push("nvtx-local-spgemm") @@ -331,6 +373,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) torch.cuda.nvtx.range_pop() # local-spgemm + # 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']}") torch.cuda.nvtx.range_push("nvtx-matc-coalesce") start_time(start_timer) matc = matc.coalesce() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 8c30dc1..6a388c0 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -433,7 +433,7 @@ def main(args, batches=None): # del adj_matrix # Comment when testing print(f"rank: {rank} g_loc: {g_loc}") - print(f"rank: {rank} batches_loc: {batches_loc}") + print(f"rank: {rank} batches_loc: {batches_loc}", flush=True) # do it once before timing torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm From d56855efa90645832c9fd7e551202a780893dfc9 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 5 Jan 2023 14:08:41 -0800 Subject: [PATCH 108/178] Included waits for isend/irecv asnd fixed bug on frontier_nnz_sizes --- cagnet/samplers/utils.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index ed69891..cf49b18 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -266,11 +266,11 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups for j in range(size // replication): recv_rank = rank_col + j * replication if recv_rank != q: - recv_objs.append(dist.irecv(nnz_col_ids[j], src=recv_rank).wait()) + recv_objs.append(dist.recv(nnz_col_ids[j], src=recv_rank)) else: nnz_col_ids[j] = nnz_cols else: - dist.isend(nnz_cols, dst=q) + dist.send(nnz_cols, dst=q) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-colids-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-nnzcols @@ -309,10 +309,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_indices.size(1)) nnz_count += 3 * matb_send_indices.size(1) - dist.isend(selected_rows_count, tag=0, dst=recv_rank) + rows_count_obj = dist.isend(selected_rows_count, tag=0, dst=recv_rank) # dist.isend(matb_send, tag=1, dst=recv_rank) - dist.isend(matb_send_indices, tag=1, dst=recv_rank) - dist.isend(matb_send_values, tag=2, dst=recv_rank) + matb_indices_obj = dist.isend(matb_send_indices, tag=1, dst=recv_rank) + matb_values_obj = dist.isend(matb_send_values, tag=2, dst=recv_rank) + + rows_count_obj.wait() + matb_indices_obj.wait() + matb_values_obj.wait() else: # matb_select_recv = matb_send.clone() matb_recv_indices = matb_send_indices.clone() @@ -543,10 +547,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], bins=p.size(0)).int() frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0)).fill_(0) - frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], next_frontier._indices()[0,:].int()) - frontier_nnz_sizes = torch.div(frontier_nnz_sizes, - torch.arange(0, p.size(0), device=torch.device("cuda:0")), - rounding_mode="floor").int() + ones = torch.cuda.IntTensor(next_frontier._indices()[0,:].size(0)).fill_(1) + frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) zero_count = (next_frontier._indices()[0,:] == 0).nonzero().size(0) frontier_nnz_sizes[0] = zero_count From 319750deef9127344a4c257e2edd89bc49648b92 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 10 Jan 2023 14:38:00 -0800 Subject: [PATCH 109/178] Extracted rowselect coo, correct on cora --- cagnet/samplers/utils.py | 11 ++-- examples/gcn_15d.py | 10 +--- sparse-extension/sparse_coo_tensor.cu | 82 ++++++++++++++++++++++----- tests/test_sage.py | 8 ++- 4 files changed, 83 insertions(+), 28 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index cf49b18..21170cc 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -84,7 +84,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, start_time(start_timer) matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ - mata_chunk_values, matb_recv_indices, \ + mata_chunk_values.double(), matb_recv_indices, \ matb_recv_values, mata.size(0), \ chunk_col_size, matb.size(1), coalesced=True) @@ -279,13 +279,17 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_time(start_timer) nnz_count = 0 if rank == q: + nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) + rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ + matb._indices().size(1), size // replication) for j in range(size // replication): # start_time(start_inner_timer) recv_rank = rank_col + j * replication nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] - rowselect_coo_gpu(nnz_col_ids[j], matb._indices()[0,:], nnz_row_mask, \ - nnz_col_ids[j].size(0), matb._indices().size(1)) + # rowselect_coo_gpu(nnz_col_ids[j], matb._indices()[0,:], nnz_row_mask, \ + # nnz_col_ids[j].size(0), matb._indices().size(1)) + # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowselect-{name}") # matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) @@ -544,7 +548,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) torch.cuda.nvtx.range_push("nvtx-pre-loop") - # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,:], bins=p.size(0)).int() frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0)).fill_(0) ones = torch.cuda.IntTensor(next_frontier._indices()[0,:].size(0)).fill_(1) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 6a388c0..f16708e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -492,7 +492,6 @@ def main(args, batches=None): # adj_matrices[i][j] -- mb i layer j rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] - print(f"before reindexing", flush=True) for i in range(args.n_layers): for j in range(rank_n_bulkmb): row_select_min = j * args.batch_size @@ -517,7 +516,6 @@ def main(args, batches=None): (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) adj_matrices[j][i] = adj_matrix_sample - print(f"after reindexing", flush=True) # print(f"sample: {adj_matrices}") # adj_matrix = None # Uncomment bottom for testing adj_matrix = adj_matrix.cuda() @@ -529,12 +527,8 @@ def main(args, batches=None): adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code - print(f"current_frontier: {current_frontier}") - - print(f"next_frontier[20]: {next_frontier[(next_frontier[:,-1] == 20).nonzero().squeeze(), :]}") - print(f"next_frontier: {next_frontier}") - print(f"adj_matrices: {adj_matrices}") - return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + return current_frontier, next_frontier, adj_matrices, adj_matrix + # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups return # create GCN model diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index ee57ca0..492c8f7 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -835,9 +835,9 @@ __global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_va for (int i = id; i < dart_count; i += stride) { int vtx = binary_searchf(ps_p_values, dart_values[i], 0, nnz - 1); - if (vtx < 0 || vtx >= nnz) { - printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); - } + // if (vtx < 0 || vtx >= nnz) { + // printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); + // } atomicAdd(&h_values[vtx], 1); } } @@ -1021,35 +1021,87 @@ void scatteri_add_gpu(const at::Tensor& src, const at::Tensor& indices, const at CHECK_ERROR("scatter add ints error") } -__global__ void RowSelectCoo(long *nnz_cols, long *row_ids, bool *mask, int nnz_col_count, int row_count) { +// // 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; - 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; + 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(const at::Tensor& nnz_cols, const at::Tensor& rows, const at::Tensor& mask, - int nnz_col_count, int row_count) { +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(row_count / ((float) BLOCK_SIZE)); + int BLOCK_COUNT = std::ceil(mask_len / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - if (nnz_col_count == 0 || row_count == 0) { + if (mask_len == 0) { return; } - RowSelectCoo<<>>(nnz_cols.data(), + 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_col_count, - row_count); + nnz_cols_counts.data(), + row_count, + proc_count); CHECK_ERROR("rowselect coo error") } diff --git a/tests/test_sage.py b/tests/test_sage.py index bf7c959..344ee67 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -48,11 +48,13 @@ def main(args): agg_neighbors = torch.IntTensor(agg_neighbors) agg_neighbors = agg_neighbors.unique() + print(f"current_frontier: {current_frontier}") + print(f"next_frontier: {next_frontier}") # 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: + if vtx not in agg_neighbors and vtx not in current_frontier: vertices_match = False print(f"vtx {vtx} not in aggregated neighborhood") @@ -122,6 +124,10 @@ def main(args): help='Number of minibatches to sample in bulk') parser.add_argument('--n-darts', type=int, default=-1, help='Number of darts to throw per minibatch in LADIES sampling') + parser.add_argument('--timing', action="store_true", + help='whether to turn on timers') + parser.add_argument('--baseline', action="store_true", + help='whether to avoid col selection for baseline comparison') args = parser.parse_args() print(args) From e69b75700e65f96454decea718cd79710858e3dd Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 12 Jan 2023 14:42:11 -0800 Subject: [PATCH 110/178] Uses nccl batch_isend_irecv --- cagnet/samplers/utils.py | 70 ++++++++++++++++++++++++++-------------- examples/gcn_15d.py | 3 +- 2 files changed, 48 insertions(+), 25 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 21170cc..f0becd8 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -282,6 +282,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ matb._indices().size(1), size // replication) + + send_ops = [None] * 3 * ((size // replication) - 1) + send_idx = 0 for j in range(size // replication): # start_time(start_inner_timer) recv_rank = rank_col + j * replication @@ -313,18 +316,26 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_indices.size(1)) nnz_count += 3 * matb_send_indices.size(1) - rows_count_obj = dist.isend(selected_rows_count, tag=0, dst=recv_rank) - # dist.isend(matb_send, tag=1, dst=recv_rank) - matb_indices_obj = dist.isend(matb_send_indices, tag=1, dst=recv_rank) - matb_values_obj = dist.isend(matb_send_values, tag=2, dst=recv_rank) - - rows_count_obj.wait() - matb_indices_obj.wait() - matb_values_obj.wait() + # dist.isend(selected_rows_count, tag=0, dst=recv_rank) + # # dist.isend(matb_send, tag=1, dst=recv_rank) + # dist.isend(matb_send_indices, tag=1, dst=recv_rank) + # dist.isend(matb_send_values, tag=2, dst=recv_rank) + + # send_ops.extend([dist.P2POp(dist.isend, selected_rows_count, recv_rank, tag=0), \ + # dist.P2POp(dist.isend, matb_send_indices, recv_rank, tag=1), \ + # dist.P2POp(dist.isend, matb_send_values, recv_rank, tag=2)]) + send_ops[send_idx] = dist.P2POp(dist.isend, selected_rows_count.clone(), recv_rank, tag=0) + send_ops[send_idx + 1] = dist.P2POp(dist.isend, matb_send_indices.clone(), recv_rank, tag=1) + send_ops[send_idx + 2] = dist.P2POp(dist.isend, matb_send_values.clone(), 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() + + reqs = dist.batch_isend_irecv(send_ops) + for req in reqs: + req.wait() else: # if f"spgemm-rowselect-{name}" not in timing_dict: # timing_dict[f"spgemm-rowselect-{name}"] = [] @@ -335,35 +346,46 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # if f"spgemm-rowval-{name}" not in timing_dict: # timing_dict[f"spgemm-rowval-{name}"] = [] + # if f"spgemm-rowmisc-{name}" not in timing_dict: + # timing_dict[f"spgemm-rowmisc-{name}"] = [] + selected_rows_count_recv = torch.cuda.IntTensor(1) - dist.irecv(selected_rows_count_recv, tag=0, src=q).wait() + dist.recv(selected_rows_count_recv, tag=0, src=q) # 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()) - ind_obj = dist.irecv(matb_recv_indices, tag=1, src=q) + # dist.irecv(matb_recv_indices, tag=1, src=q) matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) - val_obj = dist.irecv(matb_recv_values, tag=2, src=q) + # dist.irecv(matb_recv_values, tag=2, src=q) - ind_obj.wait() - val_obj.wait() + recv_ops = [dist.P2POp(dist.irecv, matb_recv_indices, q, tag=1), \ + dist.P2POp(dist.irecv, matb_recv_values, q, tag=2)] - nnz_count += 3 * matb_recv_indices.size(1) + 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_timer, stop_timer, timing_dict, f"spgemm-send-rowdata-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-rowcounts - # timing_dict[f"spgemm-send-rownnz-{name}"].append(nnz_count) - torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") - 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, :] + timing_dict[f"spgemm-send-rownnz-{name}"].append(nnz_count) + if i == stages - 1 and rank == q: + print(f"rowdata-time: {timing_dict['spgemm-send-rowdata-prob']}") + print(f"rowdata-nnz: {timing_dict['spgemm-send-rownnz-prob']}") + + # torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") + # 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) + # # 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) + # torch.cuda.nvtx.range_pop() # spgemm-castrecv torch.cuda.nvtx.range_push("nvtx-local-spgemm") start_time(start_timer) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index f16708e..f3b081d 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -465,6 +465,7 @@ def main(args, batches=None): col_groups, args.timing) elif args.sample_method == "sage": torch.manual_seed(rank_col) + print("first (warmup) run") current_frontier, next_frontier, adj_matrices = \ sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -473,7 +474,7 @@ def main(args, batches=None): sa_recv_buff, rank, size, row_groups, col_groups, args.timing, args.baseline) - print() + print("second run") torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") # torch.manual_seed(rank_col) From cf245f4382a9c9f42113d61723a8996619d31231 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sun, 22 Jan 2023 15:30:33 -0800 Subject: [PATCH 111/178] blocking vs nonblocking nccl microbenchmark --- .gitignore | 1 + cagnet/samplers/utils.py | 60 +++---- tests/test_nccl.cu | 359 ++++++++++++++++++++++++++++++--------- tests/test_nccl_p2p.py | 322 +++++++++++++++++++++++++++++++---- 4 files changed, 594 insertions(+), 148 deletions(-) diff --git a/.gitignore b/.gitignore index caa8558..c852688 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,7 @@ reddit.py tests/data tests/examples tests/LADIES/ +tests/test_nccl tests/pytorch_geometric/ tests/torch-quiver/ tests/*.out diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index f0becd8..a6ae6bf 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -283,7 +283,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ matb._indices().size(1), size // replication) - send_ops = [None] * 3 * ((size // replication) - 1) + recv_proc_count = ((size // replication) - 1) + send_ops = [None] * 3 * recv_proc_count send_idx = 0 for j in range(size // replication): # start_time(start_inner_timer) @@ -316,26 +317,24 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_indices.size(1)) nnz_count += 3 * matb_send_indices.size(1) - # dist.isend(selected_rows_count, tag=0, dst=recv_rank) - # # dist.isend(matb_send, tag=1, dst=recv_rank) - # dist.isend(matb_send_indices, tag=1, dst=recv_rank) - # dist.isend(matb_send_values, tag=2, dst=recv_rank) - - # send_ops.extend([dist.P2POp(dist.isend, selected_rows_count, recv_rank, tag=0), \ - # dist.P2POp(dist.isend, matb_send_indices, recv_rank, tag=1), \ - # dist.P2POp(dist.isend, matb_send_values, recv_rank, tag=2)]) - send_ops[send_idx] = dist.P2POp(dist.isend, selected_rows_count.clone(), recv_rank, tag=0) - send_ops[send_idx + 1] = dist.P2POp(dist.isend, matb_send_indices.clone(), recv_rank, tag=1) - send_ops[send_idx + 2] = dist.P2POp(dist.isend, matb_send_values.clone(), recv_rank, tag=2) - send_idx += 3 + dist.send(selected_rows_count, tag=0, dst=recv_rank) + # dist.isend(matb_send, tag=1, 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() - reqs = dist.batch_isend_irecv(send_ops) - for req in reqs: - req.wait() + # if len(send_ops) > 0: + # reqs = dist.batch_isend_irecv(send_ops) + # for req in reqs: + # req.wait() else: # if f"spgemm-rowselect-{name}" not in timing_dict: # timing_dict[f"spgemm-rowselect-{name}"] = [] @@ -354,26 +353,24 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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.irecv(matb_recv_indices, tag=1, src=q) + dist.recv(matb_recv_indices, tag=1, src=q) matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) - # dist.irecv(matb_recv_values, tag=2, src=q) + 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)] + # 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) + # reqs = dist.batch_isend_irecv(recv_ops) + # for req in reqs: + # req.wait() + nnz_count += 3 * matb_recv_indices.size(1) + torch.cuda.synchronize() stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowdata-{name}", barrier=True) torch.cuda.nvtx.range_pop() # send-rowcounts timing_dict[f"spgemm-send-rownnz-{name}"].append(nnz_count) - if i == stages - 1 and rank == q: - print(f"rowdata-time: {timing_dict['spgemm-send-rowdata-prob']}") - print(f"rowdata-nnz: {timing_dict['spgemm-send-rownnz-prob']}") # torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") # start_time(start_timer) @@ -512,14 +509,19 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication stop_timer = torch.cuda.Event(enable_timing=True) # TODO: assume n_layers=1 for now - start_time(start_timer) + # start_time(start_timer) + start_timer.record() torch.cuda.nvtx.range_push("nvtx-probability-spgemm") p_num_indices, p_num_values = dist_saspgemm15D( numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, sa_recv_buff, timing_dict) torch.cuda.nvtx.range_pop() - timing_dict["probability-spgemm"].append(stop_time(start_timer, stop_timer)) + 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) torch.cuda.nvtx.range_push("nvtx-compute-p") diff --git a/tests/test_nccl.cu b/tests/test_nccl.cu index fb9faa4..1b05cbb 100644 --- a/tests/test_nccl.cu +++ b/tests/test_nccl.cu @@ -1,9 +1,22 @@ -#include #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; \ @@ -24,116 +37,296 @@ } \ } 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; +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; +} + - int n; - int ngpus; - n = atoi(argv[1]); - ngpus = atoi(argv[2]); - n *= 1000; +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 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); +int main(int argc, char* argv[]) +{ + int size = 1*1024*1024*1024; - 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; - } + int myRank, nRanks, localRank = 0; - // 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)); - } + //initializing MPI + MPICHECK(MPI_Init(&argc, &argv)); + MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank)); + MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks)); - CUDACHECK(cudaMemset(recvbuff[i], 0, n * sizeof(float))); - CUDACHECK(cudaStreamCreate(s+i)); - CUDACHECK(cudaEventCreate(&start[i])); - CUDACHECK(cudaEventCreate(&stop[i])); + //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: @@ -93,6 +93,256 @@ def test_nccl(args): # 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 + 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() + torch.cuda.nvtx.range_pop() + + # print(f"batch send") + # 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) + # dist.recv(recv_tensor, src=0) + # 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() + 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) + # dist.barrier() + torch.cuda.nvtx.range_pop() + torch.cuda.profiler.cudart().cudaProfilerStop() if __name__ == '__main__': parser = argparse.ArgumentParser(description='NCCL P2P test') From 54f0ac2c66520053ca58cbc83feec7cb243f3781 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 27 Jan 2023 10:51:30 -0800 Subject: [PATCH 112/178] Added new job file and more timers --- cagnet/samplers/utils.py | 51 ++++++++++++++-------------------------- examples/job_sage | 20 ---------------- examples/job_sage_p36 | 49 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 53 deletions(-) delete mode 100755 examples/job_sage create mode 100755 examples/job_sage_p36 diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index a6ae6bf..427da5e 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -279,46 +279,31 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_time(start_timer) nnz_count = 0 if rank == q: + start_time(start_inner_timer) nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ matb._indices().size(1), size // replication) - recv_proc_count = ((size // replication) - 1) - send_ops = [None] * 3 * recv_proc_count - send_idx = 0 + # recv_proc_count = ((size // replication) - 1) + # send_ops = [None] * 3 * recv_proc_count + # send_idx = 0 + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") + + start_time(start_inner_timer) for j in range(size // replication): - # start_time(start_inner_timer) recv_rank = rank_col + j * replication nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] # rowselect_coo_gpu(nnz_col_ids[j], matb._indices()[0,:], nnz_row_mask, \ # nnz_col_ids[j].size(0), matb._indices().size(1)) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowselect-{name}") - - # matb_send_indices = torch.cuda.LongTensor(2, nnz_row_mask.nonzero().size(1)) - # torch.cuda.synchronize() - - # start_time(start_inner_timer) matb_send_indices = matb._indices()[:, nnz_row_mask] - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowind-{name}") - - # start_time(start_inner_timer) matb_send_values = matb._values()[nnz_row_mask] - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowval-{name}") - # start_time(start_inner_timer) - # # matb_send_values = matb_send_values.unsqueeze(dim=0) - # # matb_send_indices = matb_send_indices.float() - # # matb_send = torch.cat((matb_send_indices, matb_send_values), dim=0) - # stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-rowmisc-{name}") - if recv_rank != q: - # selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send.size(1)) 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.isend(matb_send, tag=1, dst=recv_rank) dist.send(matb_send_indices, tag=1, dst=recv_rank) dist.send(matb_send_values, tag=2, dst=recv_rank) @@ -330,24 +315,20 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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: # reqs = dist.batch_isend_irecv(send_ops) # for req in reqs: # req.wait() else: - # if f"spgemm-rowselect-{name}" not in timing_dict: - # timing_dict[f"spgemm-rowselect-{name}"] = [] - - # if f"spgemm-rowind-{name}" not in timing_dict: - # timing_dict[f"spgemm-rowind-{name}"] = [] - - # if f"spgemm-rowval-{name}" not in timing_dict: - # timing_dict[f"spgemm-rowval-{name}"] = [] - - # if f"spgemm-rowmisc-{name}" not in timing_dict: - # timing_dict[f"spgemm-rowmisc-{name}"] = [] + if f"spgemm-send-rowsel-{name}" not in timing_dict: + timing_dict[f"spgemm-send-rowsel-{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) dist.recv(selected_rows_count_recv, tag=0, src=q) @@ -365,11 +346,15 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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) torch.cuda.nvtx.range_pop() # send-rowcounts + 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) # torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") diff --git a/examples/job_sage b/examples/job_sage deleted file mode 100755 index 32821e0..0000000 --- a/examples/job_sage +++ /dev/null @@ -1,20 +0,0 @@ -#!/bin/bash -#SBATCH --qos=regular -#SBATCH --time=00:30:00 -#SBATCH --nodes=9 -#SBATCH --gpus=36 -#SBATCH --ntasks-per-node=4 -#SBATCH --constraint=gpu -#SBATCH --account=m1982_g - -module load pytorch -export PYTHONPATH="/global/homes/a/alokt/CAGNET:$PYTHONPATH" -export PYTHONPATH="/global/homes/a/alokt/CAGNET/sparse-extension/build/lib/python3.9/site-packages/sparse_coo_tensor_cpp-0.0.0-py3.9-linux-x86_64.egg:$PYTHONPATH" - -hostname=$(scontrol show hostnames | head -n 1) -echo $hostname -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_coogeam.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_coogeam.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_coogeam.out diff --git a/examples/job_sage_p36 b/examples/job_sage_p36 new file mode 100755 index 0000000..d1c62d4 --- /dev/null +++ b/examples/job_sage_p36 @@ -0,0 +1,49 @@ +#!/bin/bash +#SBATCH --qos=regular +#SBATCH --time=01:00:00 +#SBATCH --nodes=9 +#SBATCH --gpus=36 +#SBATCH --ntasks-per-node=4 +#SBATCH --constraint=gpu +#SBATCH --account=m1982_g + +hostname=$(scontrol show hostnames | head -n 1) +echo $hostname +# Protein +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out + +# Amazon +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out + +# Reddit +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out From e116046eff07e03f448862c01b6bbc2fbe4ea235 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 31 Jan 2023 17:57:01 -0800 Subject: [PATCH 113/178] Partway through COO to CSR conversion --- cagnet/samplers/sage.py | 4 + cagnet/samplers/utils.py | 105 ++++++++++++++++++-------- sparse-extension/sparse_coo_tensor.cu | 33 ++++++++ 3 files changed, 112 insertions(+), 30 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 738f44f..176e6e6 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -89,10 +89,14 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches._values(), size=(mb_count * nnz, node_count_total)) + batches_expand = batches_expand.to_sparse_csr() + adj_matrix = adj_matrix.to_sparse_csr() p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "sage", timing_arg) + adj_matrix = adj_matrix.to_sparse_coo() + p = p.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 427da5e..4cc6da3 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -197,7 +197,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-matc-inst") 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() + if mata.layout == torch.sparse_coo: + 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=matb.size()) rank_c = rank // replication rank_col = rank % replication stages = size // (replication ** 2) @@ -217,18 +224,32 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) + 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] + mata_chunk_rowcount = torch.segment_reduce(mask.float(), "sum", offsets=mata.crow_indices()) + mata_chunk_crows = torch.cuda.LongTensor(mata.size(0) + 1) + mata_chunk_crows[1:] = torch.cumsum(mata_chunk_crows, dim=0).long() + mata_chunk_crows[0] = 0 - 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) torch.cuda.nvtx.range_pop() # preproc-local # Determine number of nonzero columns in chunk of mata torch.cuda.nvtx.range_push("nvtx-unique") start_time(start_timer) - nnz_cols = torch.unique(mata_chunk_indices[1,:]) + 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) torch.cuda.nvtx.range_pop() # unique @@ -281,9 +302,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups if rank == q: start_time(start_inner_timer) nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) - rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ - matb._indices().size(1), size // replication) - + if matb.layout == torch.sparse_coo: + rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ + matb._indices().size(1), size // replication) # recv_proc_count = ((size // replication) - 1) # send_ops = [None] * 3 * recv_proc_count # send_idx = 0 @@ -294,27 +315,51 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups recv_rank = rank_col + j * replication nnz_row_mask = nnz_row_masks[(j * matb._indices().size(1)):((j + 1) * matb._indices().size(1))] - # rowselect_coo_gpu(nnz_col_ids[j], matb._indices()[0,:], nnz_row_mask, \ - # nnz_col_ids[j].size(0), matb._indices().size(1)) - - matb_send_indices = matb._indices()[:, nnz_row_mask] - matb_send_values = matb._values()[nnz_row_mask] - - 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() + if matb.layout == torch.sparse_csr: + rowselect_csr_gpu(nnz_col_ids[j], matb.crow_indices(), nnz_row_mask, \ + nnz_col_ids[j].size(0), matb._nnz()) + + matb_send_crows = matb.crow_indices() + matb_send_cols = matb.col_indices()[nnz_row_mask] + matb_send_values = matb.values()[nnz_row_mask] + + if recv_rank != q: + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_cols.size(1)) + nnz_count += 3 * matb._nnz().item() + 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() + 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)) + + matb_send_indices = matb._indices()[:, nnz_row_mask] + matb_send_values = matb._values()[nnz_row_mask] + + 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: diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 492c8f7..9859c2d 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1105,6 +1105,39 @@ void rowselect_coo_gpu(std::vector nnz_cols, const at::Tensor& rows, 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 = nnz_cols[vtx + 1] - nnz_cols[vtx]; + for (int j = 0; j < degree; j++) { + mask[row_offsets[vtx] + j] = true; + } + } +} + +void rowselect_coo_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") +} + 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"); From bcdffd98359c2245ffe2682b18a8cd475d165be7 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 3 Feb 2023 02:13:32 -0800 Subject: [PATCH 114/178] Temporary CSR sa-spgemm implementation --- cagnet/samplers/sage.py | 3 +- cagnet/samplers/utils.py | 166 +++++++++++++++++++------- examples/gcn_15d.py | 2 +- examples/job_sage_p36 | 76 ++++++------ sparse-extension/sparse_coo_tensor.cu | 8 +- tests/test_sage.py | 2 - 6 files changed, 169 insertions(+), 88 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 176e6e6..25a4cdb 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -96,7 +96,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, sa_masks, sa_recv_buff, timing_dict, "sage", timing_arg) adj_matrix = adj_matrix.to_sparse_coo() - p = p.to_sparse_coo() + if p.layout == torch.sparse_csr: + p = p.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 4cc6da3..eb99234 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -9,7 +9,7 @@ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ - sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu + rowselect_csr_gpu, sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu timing = True @@ -204,7 +204,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matc = torch.sparse_csr_tensor(matc_crow_indices, \ torch.cuda.LongTensor(0), \ torch.cuda.FloatTensor(0), - size=matb.size()) + size=(mata.size(0), matb.size(1))) rank_c = rank // replication rank_col = rank % replication stages = size // (replication ** 2) @@ -225,19 +225,30 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) + 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) + 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] - mata_chunk_rowcount = torch.segment_reduce(mask.float(), "sum", offsets=mata.crow_indices()) + + # # 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_crows, dim=0).long() + mata_chunk_crows[1:] = torch.cumsum(mata_chunk_rowcount, dim=0).long() mata_chunk_crows[0] = 0 stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-preproc-local-{name}", barrier=True) @@ -300,51 +311,74 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_time(start_timer) nnz_count = 0 if rank == q: - start_time(start_inner_timer) nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) if matb.layout == torch.sparse_coo: + start_time(start_inner_timer) rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ 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] * 3 * recv_proc_count # send_idx = 0 - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") - start_time(start_inner_timer) + # 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._indices().size(1)):((j + 1) * matb._indices().size(1))] + nnz_row_mask = nnz_row_masks[(j * matb._nnz()):((j + 1) * matb._nnz())] if matb.layout == torch.sparse_csr: + start_time(start_inner_timer) rowselect_csr_gpu(nnz_col_ids[j], matb.crow_indices(), nnz_row_mask, \ nnz_col_ids[j].size(0), matb._nnz()) - - matb_send_crows = matb.crow_indices() + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") + + matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) + lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] + if nnz_row_mask.any(): + matb_rows_sum[nnz_col_ids[j]] = lengths[nnz_col_ids[j]] + + # # torch 1.13 supports offsets-based reduction + start_time(start_inner_timer) + # if nnz_row_mask.any(): + # lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] + # matb_rows_sum = torch.segment_reduce(nnz_row_mask.float(), + # "sum", + # lengths=lengths) + # # matb_rows_sum = torch.segment_reduce(nnz_row_mask.float(), + # # "sum", + # # offsets=matb.crow_indices()) + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-segred-{name}") + start_time(start_inner_timer) + matb_send_crows = torch.cuda.LongTensor(matb.crow_indices().size(0)) + matb_send_crows[1:] = torch.cumsum(matb_rows_sum, dim=0) + matb_send_crows[0] = 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(1)) - nnz_count += 3 * matb._nnz().item() + selected_rows_count = torch.cuda.IntTensor(1).fill_(matb_send_cols.size(0)) + nnz_count += 2 * matb_send_cols.size(0) + matb_send_crows.size(0) 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 + dist.send(matb_send_crows, 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) else: - # matb_select_recv = matb_send.clone() - matb_recv_indices = matb_send_indices.clone() + matb_recv_crows = matb_send_crows.clone() + matb_recv_cols = matb_send_cols.clone() matb_recv_values = matb_send_values.clone() + 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) @@ -360,7 +394,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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}") + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-calls-{name}") # if len(send_ops) > 0: # reqs = dist.batch_isend_irecv(send_ops) @@ -370,27 +404,49 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}"] = [] + if matb.layout == torch.sparse_csr: + if f"spgemm-send-segred-{name}" not in timing_dict: + timing_dict[f"spgemm-send-segred-{name}"] = [] + start_time(start_inner_timer) selected_rows_count_recv = torch.cuda.IntTensor(1) dist.recv(selected_rows_count_recv, tag=0, src=q) - # 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) + 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_cols = torch.cuda.LongTensor(selected_rows_count_recv.item()) + dist.recv(matb_recv_cols, tag=2, src=q) + + matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) + dist.recv(matb_recv_values, tag=3, src=q) + + print(f"nnz: {matb_recv_cols.size(0)} rows: {matb_recv_crows.size(0)}") + nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_crows.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) + 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)] + # 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) + # 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() @@ -417,14 +473,36 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-local-spgemm") start_time(start_timer) if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: - 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 mata.layout == torch.sparse_csr: + 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))) + + + # matc += torch.mm(mata_recv, matb_recv) + + mata_recv = mata_recv.to_sparse_coo().coalesce() + matb_recv = matb_recv.to_sparse_coo().coalesce() + 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) + 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.to_sparse_coo() + matc += matc_chunk + 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 else: pass stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) @@ -433,6 +511,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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']}") + if matc.layout == torch.sparse_csr: + matc = matc.to_sparse_coo() torch.cuda.nvtx.range_push("nvtx-matc-coalesce") start_time(start_timer) matc = matc.coalesce() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index f3b081d..2ee448e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -474,7 +474,7 @@ def main(args, batches=None): sa_recv_buff, rank, size, row_groups, col_groups, args.timing, args.baseline) - print("second run") + print("second run", flush=True) torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") # torch.manual_seed(rank_col) diff --git a/examples/job_sage_p36 b/examples/job_sage_p36 index d1c62d4..31ed3f1 100755 --- a/examples/job_sage_p36 +++ b/examples/job_sage_p36 @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --qos=regular -#SBATCH --time=01:00:00 +#SBATCH --time=00:10:00 #SBATCH --nodes=9 #SBATCH --gpus=36 #SBATCH --ntasks-per-node=4 @@ -10,40 +10,40 @@ hostname=$(scontrol show hostnames | head -n 1) echo $hostname # Protein -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out - -# Amazon -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out - -# Reddit -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_csr_nosegred.out + +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_csr_nosegred.out + +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out + +# # Amazon +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out +# +# # Reddit +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out +# +# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 9859c2d..f4e84c0 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1112,14 +1112,15 @@ __global__ void RowSelectCsr(long *nnz_cols, long *row_offsets, bool *mask, int for (int i = id; i < nnz_col_count; i += stride) { long vtx = nnz_cols[i]; - long degree = nnz_cols[vtx + 1] - nnz_cols[vtx]; + 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_coo_gpu(const at::Tensor& nnz_cols, const at::Tensor& row_offsets, const at::Tensor& mask, +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) { @@ -1154,6 +1155,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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, "Row selection for sparsity-aware spgemm"); + 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"); } diff --git a/tests/test_sage.py b/tests/test_sage.py index 344ee67..778454e 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -48,8 +48,6 @@ def main(args): agg_neighbors = torch.IntTensor(agg_neighbors) agg_neighbors = agg_neighbors.unique() - print(f"current_frontier: {current_frontier}") - print(f"next_frontier: {next_frontier}") # Verify each vertex in next_frontier exists in aggregated neighorhood vertices_match = True for v in next_frontier: From 805bfde2ccd749dc5fefc89d48aec1ae4ad17714 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 3 Feb 2023 12:14:12 -0800 Subject: [PATCH 115/178] sends lengths instead of csr offsets --- cagnet/samplers/sage.py | 10 ++-- cagnet/samplers/utils.py | 62 ++++++++++++----------- examples/job_sage_p36 | 4 +- tests/test_nccl_p2p.py | 106 +++++++++++++++++++++------------------ 4 files changed, 97 insertions(+), 85 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 25a4cdb..52d5c2a 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -89,15 +89,15 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches._values(), size=(mb_count * nnz, node_count_total)) - batches_expand = batches_expand.to_sparse_csr() - adj_matrix = adj_matrix.to_sparse_csr() + # batches_expand = batches_expand.to_sparse_csr() + # adj_matrix = adj_matrix.to_sparse_csr() p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "sage", timing_arg) - adj_matrix = adj_matrix.to_sparse_coo() - if p.layout == torch.sparse_csr: - p = p.to_sparse_coo() + # adj_matrix = adj_matrix.to_sparse_coo() + # if p.layout == torch.sparse_csr: + # p = p.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index eb99234..bd61de4 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -332,26 +332,20 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_col_ids[j].size(0), matb._nnz()) stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") - matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) - lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] - if nnz_row_mask.any(): - matb_rows_sum[nnz_col_ids[j]] = lengths[nnz_col_ids[j]] - - # # torch 1.13 supports offsets-based reduction - start_time(start_inner_timer) + # matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) + # lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] # if nnz_row_mask.any(): - # lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] - # matb_rows_sum = torch.segment_reduce(nnz_row_mask.float(), - # "sum", - # lengths=lengths) - # # matb_rows_sum = torch.segment_reduce(nnz_row_mask.float(), - # # "sum", - # # offsets=matb.crow_indices()) - stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-segred-{name}") + # matb_rows_sum[nnz_col_ids[j]] = lengths[nnz_col_ids[j]] + 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) + start_time(start_inner_timer) - matb_send_crows = torch.cuda.LongTensor(matb.crow_indices().size(0)) - matb_send_crows[1:] = torch.cumsum(matb_rows_sum, dim=0) - matb_send_crows[0] = 0 + # matb_send_crows = torch.cuda.LongTensor(matb.crow_indices().size(0)) + # matb_send_crows[1:] = torch.cumsum(matb_rows_sum, dim=0) + # matb_send_crows[0] = 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}") @@ -359,13 +353,20 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) + matb_send_crows.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_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) else: - matb_recv_crows = matb_send_crows.clone() + # matb_recv_crows = matb_send_crows.clone() + matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) + if nnz_row_mask.any(): + matb_rows_sum[nnz_col_ids[j]] = row_lengths[nnz_col_ids[j]] + 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[0] = 0 matb_recv_cols = matb_send_cols.clone() matb_recv_values = matb_send_values.clone() stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-calls-{name}") @@ -410,17 +411,16 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups if f"spgemm-send-calls-{name}" not in timing_dict: timing_dict[f"spgemm-send-calls-{name}"] = [] - if matb.layout == torch.sparse_csr: - if f"spgemm-send-segred-{name}" not in timing_dict: - timing_dict[f"spgemm-send-segred-{name}"] = [] - start_time(start_inner_timer) selected_rows_count_recv = torch.cuda.IntTensor(1) 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_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)) + dist.recv(matb_recv_lengths, tag=1, src=q) matb_recv_cols = torch.cuda.LongTensor(selected_rows_count_recv.item()) dist.recv(matb_recv_cols, tag=2, src=q) @@ -428,7 +428,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) dist.recv(matb_recv_values, tag=3, src=q) - print(f"nnz: {matb_recv_cols.size(0)} rows: {matb_recv_crows.size(0)}") + 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, dim=0) + matb_recv_crows[0] = 0 + nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_crows.size(0) elif matb.layout == torch.sparse_coo: diff --git a/examples/job_sage_p36 b/examples/job_sage_p36 index 31ed3f1..b50d842 100755 --- a/examples/job_sage_p36 +++ b/examples/job_sage_p36 @@ -12,9 +12,9 @@ echo $hostname # Protein # srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_csr_nosegred.out +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_coo_sendlens.out -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_csr_nosegred.out +srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_coo_sendlens.out # srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out # diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index 8a235ba..7825fde 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -217,6 +217,7 @@ def test_nccl(args): # 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() @@ -248,65 +249,70 @@ def test_nccl(args): dist.barrier() torch.cuda.nvtx.range_pop() - # print(f"batch send") - # 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)) + print(f"batch send") + 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)) - # if rank == 0: - # if r == run_count - 1: - # start_time(start_timer) + send_ops = [] + for i in range(size - 1): + send_ops.append(dist.P2POp(dist.isend, tensors[i], i + 1)) - # 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 rank == 0: + if r == run_count - 1: + start_time(start_timer) - # 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() - # torch.cuda.nvtx.range_pop() + 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() - 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)) + 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: - tensors = None + 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() + torch.cuda.nvtx.range_pop() - recv_tensor = torch.cuda.FloatTensor(buffer_size) + # 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 - if rank == 0 and r == run_count - 1: - start_time(start_timer) + # recv_tensor = torch.cuda.FloatTensor(buffer_size) - dist.scatter(recv_tensor, tensors, src=0) - torch.cuda.synchronize() + # if rank == 0 and r == run_count - 1: + # start_time(start_timer) - 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() - torch.cuda.nvtx.range_pop() + # 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() + # torch.cuda.nvtx.range_pop() # print(f"nonblocking isends") # torch.cuda.nvtx.range_push(f"nvtx-nonblockisend-rank{rank}") From 641690af58e275df8030d6719fc98d4c1eb90123 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 6 Feb 2023 19:59:37 -0800 Subject: [PATCH 116/178] Added spgemm_gpu() wrapper for CSR --- cagnet/samplers/sage.py | 10 +++--- cagnet/samplers/utils.py | 48 ++++++++++++++++++------- examples/gcn_15d.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 51 +++++++++------------------ 4 files changed, 58 insertions(+), 53 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 52d5c2a..25a4cdb 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -89,15 +89,15 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches._values(), size=(mb_count * nnz, node_count_total)) - # batches_expand = batches_expand.to_sparse_csr() - # adj_matrix = adj_matrix.to_sparse_csr() + batches_expand = batches_expand.to_sparse_csr() + adj_matrix = adj_matrix.to_sparse_csr() p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "sage", timing_arg) - # adj_matrix = adj_matrix.to_sparse_coo() - # if p.layout == torch.sparse_csr: - # p = p.to_sparse_coo() + adj_matrix = adj_matrix.to_sparse_coo() + if p.layout == torch.sparse_csr: + p = p.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index bd61de4..90a162e 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -486,20 +486,42 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups size=(chunk_col_size, matb.size(1))) - # matc += torch.mm(mata_recv, matb_recv) - - mata_recv = mata_recv.to_sparse_coo().coalesce() - matb_recv = matb_recv.to_sparse_coo().coalesce() - 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) - 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 += torch.mm(mata_recv, matb_recv) + # matc_chunk_crow, matc_chunk_cols, matc_chunk_values = spgemm_gpu(mata_chunk_crows.int(), \ + # mata_chunk_cols.int(), \ + # mata_chunk_values.float(), \ + # matb_recv_crows.int(), \ + # matb_recv_cols.int(), \ + # matb_recv_values.float(), \ + # matc.crow_indices().int(), \ + # matc.col_indices().int(), \ + # matc.values().float(), \ + # mata.size(0), \ + # chunk_col_size, \ + # matb.size(1)) + + # matc_chunk = torch.sparse_csr_tensor(matc_chunk_crow.long(), \ + # matc_chunk_cols.long(), \ + # matc_chunk_values, + # size=(mata.size(0), matb.size(1))) + + # matc += matc_chunk + # matc = torch.sparse_csr_tensor(matc.crow_indices().long(), matc.col_indices().long(), + # matc.values(), matc.size()) + + + # mata_recv = mata_recv.to_sparse_coo().coalesce() + # matb_recv = matb_recv.to_sparse_coo().coalesce() + # 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) + # 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.to_sparse_coo() - matc += matc_chunk + # matc = matc.to_sparse_coo() + # matc += matc_chunk elif mata.layout == torch.sparse_coo: matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 2ee448e..9d2278e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -22,7 +22,7 @@ import cagnet.nn.functional as CAGF import ogb -from ogb.nodeproppred import PygNodePropPredDataset +# from ogb.nodeproppred import PygNodePropPredDataset import socket diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index f4e84c0..daaa0c3 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -453,12 +453,15 @@ std::vector coogeam_gpu( } std::vector spgemm_gpu( - const at::Tensor& A_rowindices, - const at::Tensor& A_colindices, + const at::Tensor& A_crow_indices, + const at::Tensor& A_col_indices, const at::Tensor& A_values, - const at::Tensor& B_rowindices, - const at::Tensor& B_colindices, + 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) { @@ -467,32 +470,12 @@ std::vector spgemm_gpu( // SpGEMM between A (n x k) and B (k x m) 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); + int32_t C_nnz = C_values.size(0); - - // Construct CSR offsets array for A and B - cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); - CHECK_CUSPARSE(cusparseXcoo2csr(handle, - A_rowindices.data(), - A_nnz, - n, - d_a_csrrows, - CUSPARSE_INDEX_BASE_ZERO)); - cudaMalloc(&d_b_csrrows, (k + 1) * sizeof(int32_t)); - - CHECK_CUSPARSE(cusparseXcoo2csr(handle, - B_rowindices.data(), - B_nnz, - k, - d_b_csrrows, - CUSPARSE_INDEX_BASE_ZERO)); - float alpha = 1; - float beta = 0; + float beta = 1; // Construct CSR matrix structs for A, B, and C (empty CSR) cusparseSpMatDescr_t matA; @@ -500,8 +483,8 @@ std::vector spgemm_gpu( n, // rows k, // cols A_nnz, // nnz - d_a_csrrows, // csrRowOffsets - A_colindices.data(), // csrColInd + A_crow_indices.data(), // csrRowOffsets + A_col_indices.data(), // csrColInd A_values.data(), // csrValues CUSPARSE_INDEX_32I, // csrRowOffsetsType CUSPARSE_INDEX_32I, // csrColIndType @@ -513,8 +496,8 @@ std::vector spgemm_gpu( k, // rows m, // cols B_nnz, // nnz - d_b_csrrows, // csrRowOffsets - B_colindices.data(), // csrColInd + B_crow_indices.data(), // csrRowOffsets + B_col_indices.data(), // csrColInd B_values.data(), // csrValues CUSPARSE_INDEX_32I, // csrRowOffsetsType CUSPARSE_INDEX_32I, // csrColIndType @@ -525,10 +508,10 @@ std::vector spgemm_gpu( CHECK_CUSPARSE(cusparseCreateCsr(&matC, n, // rows m, // cols - 0, // nnz - NULL, // csrRowOffsets - NULL, // csrColInd - NULL, // csrValues + 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, From c7c4273e9dd23bb265ab803b2bf5c4779215b64a Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 8 Feb 2023 11:59:14 -0800 Subject: [PATCH 117/178] Attempting to remove local spgemm for SAGE prob-dist --- cagnet/samplers/utils.py | 59 ++++++++++------------------------------ 1 file changed, 15 insertions(+), 44 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 90a162e..a0ce3ba 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -186,7 +186,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, return matc_nnz_indices, matc_nnz_values def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ - name, nnz_row_masks, matb_recv_buff, timing_dict=None): + name, nnz_row_masks, matb_recv_buff, timing_dict=None, alg=None): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -428,12 +428,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) dist.recv(matb_recv_values, tag=3, src=q) - 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, dim=0) - matb_recv_crows[0] = 0 + if not (name == "prob" and alg == "sage"): + 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, dim=0) + matb_recv_crows[0] = 0 nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_crows.size(0) @@ -479,7 +480,12 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-local-spgemm") start_time(start_timer) if mata_chunk_values.size(0) > 0 and matb_recv_values.size(0) > 0: - if mata.layout == torch.sparse_csr: + # Can skip local spgemm for row-selection for any alg and for SAGE's probability spgemm + if name == "prob" and alg == "sage": + 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 + elif mata.layout == torch.sparse_csr: 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(), @@ -487,41 +493,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matc += torch.mm(mata_recv, matb_recv) - # matc_chunk_crow, matc_chunk_cols, matc_chunk_values = spgemm_gpu(mata_chunk_crows.int(), \ - # mata_chunk_cols.int(), \ - # mata_chunk_values.float(), \ - # matb_recv_crows.int(), \ - # matb_recv_cols.int(), \ - # matb_recv_values.float(), \ - # matc.crow_indices().int(), \ - # matc.col_indices().int(), \ - # matc.values().float(), \ - # mata.size(0), \ - # chunk_col_size, \ - # matb.size(1)) - - # matc_chunk = torch.sparse_csr_tensor(matc_chunk_crow.long(), \ - # matc_chunk_cols.long(), \ - # matc_chunk_values, - # size=(mata.size(0), matb.size(1))) - - # matc += matc_chunk - # matc = torch.sparse_csr_tensor(matc.crow_indices().long(), matc.col_indices().long(), - # matc.values(), matc.size()) - - - # mata_recv = mata_recv.to_sparse_coo().coalesce() - # matb_recv = matb_recv.to_sparse_coo().coalesce() - # 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) - # 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.to_sparse_coo() - # matc += matc_chunk elif mata.layout == torch.sparse_coo: matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ @@ -653,7 +624,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication p_num_indices, p_num_values = dist_saspgemm15D( numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, - sa_recv_buff, timing_dict) + sa_recv_buff, timing_dict, name) torch.cuda.nvtx.range_pop() stop_timer.record() torch.cuda.synchronize() From 5556d7c533cf158b37801607979c55c59ca4ce29 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 14 Feb 2023 19:04:43 -0800 Subject: [PATCH 118/178] Converts csr2coo for very sparse storage --- .gitignore | 1 + cagnet/samplers/utils.py | 85 ++++++++++++++++++--------- examples/job_sage_p36 | 49 --------------- sparse-extension/sparse_coo_tensor.cu | 12 +++- tests/test_nccl_p2p.py | 47 ++++++++------- 5 files changed, 93 insertions(+), 101 deletions(-) delete mode 100755 examples/job_sage_p36 diff --git a/.gitignore b/.gitignore index c852688..a4255fe 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ sparse_coo_tensor_cpp.egg-info/ *.pt slurm_outputs/ job-* +job_* gcn_distr.py gcn_distr_15d.py gcn_distr_2d.py diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index a0ce3ba..f203737 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -9,7 +9,8 @@ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ - rowselect_csr_gpu, sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu + rowselect_csr_gpu, sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu, \ + sparse_csr_tensor_gpu timing = True @@ -197,7 +198,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-matc-inst") start_time(start_timer) chunk_size = math.ceil(float(mata.size(1) / (size / replication))) - if mata.layout == torch.sparse_coo: + if True or mata.layout == torch.sparse_coo: 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) @@ -332,20 +333,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_col_ids[j].size(0), matb._nnz()) stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-send-rowsel-{name}") - # matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) - # lengths = matb.crow_indices()[1:] - matb.crow_indices()[:-1] - # if nnz_row_mask.any(): - # matb_rows_sum[nnz_col_ids[j]] = lengths[nnz_col_ids[j]] + 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) - start_time(start_inner_timer) - # matb_send_crows = torch.cuda.LongTensor(matb.crow_indices().size(0)) - # matb_send_crows[1:] = torch.cumsum(matb_rows_sum, dim=0) - # matb_send_crows[0] = 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}") @@ -361,14 +355,18 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups dist.send(matb_send_values, tag=3, dst=recv_rank) else: # matb_recv_crows = matb_send_crows.clone() - matb_rows_sum = torch.cuda.LongTensor(matb.crow_indices().size(0) - 1).fill_(0) - if nnz_row_mask.any(): - matb_rows_sum[nnz_col_ids[j]] = row_lengths[nnz_col_ids[j]] - 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[0] = 0 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) + 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[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)) 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, \ @@ -435,8 +433,11 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_crows = torch.cuda.LongTensor(chunk_col_size + 1) matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, 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)) - nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_crows.size(0) + nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_lengths.size(0) elif matb.layout == torch.sparse_coo: @@ -481,18 +482,36 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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 = torch.sparse_csr_tensor(matb_recv_crows, matb_recv_cols, matb_recv_values.float(), - size=(chunk_col_size, matb.size(1))) - matc += matc_chunk - elif mata.layout == torch.sparse_csr: - mata_recv = torch.sparse_csr_tensor(mata_chunk_crows, mata_chunk_cols, mata_chunk_values.float(), + # 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: + 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)) - matb_recv = torch.sparse_csr_tensor(matb_recv_crows, matb_recv_cols, matb_recv_values.float(), - size=(chunk_col_size, matb.size(1))) + # 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}") - matc += torch.mm(mata_recv, matb_recv) + 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) + 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 + stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-spspmm-{name}") elif mata.layout == torch.sparse_coo: matc_chunk_indices, matc_chunk_values = torch_sparse.spspmm(mata_chunk_indices, \ mata_chunk_values, matb_recv_indices, \ @@ -503,17 +522,25 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk = matc_chunk.coalesce() matc += matc_chunk else: - pass + 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) torch.cuda.nvtx.range_pop() # local-spgemm # 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']}") - if matc.layout == torch.sparse_csr: - matc = matc.to_sparse_coo() torch.cuda.nvtx.range_push("nvtx-matc-coalesce") 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) torch.cuda.nvtx.range_pop() # matc-coalesce diff --git a/examples/job_sage_p36 b/examples/job_sage_p36 deleted file mode 100755 index b50d842..0000000 --- a/examples/job_sage_p36 +++ /dev/null @@ -1,49 +0,0 @@ -#!/bin/bash -#SBATCH --qos=regular -#SBATCH --time=00:10:00 -#SBATCH --nodes=9 -#SBATCH --gpus=36 -#SBATCH --ntasks-per-node=4 -#SBATCH --constraint=gpu -#SBATCH --account=m1982_g - -hostname=$(scontrol show hostnames | head -n 1) -echo $hostname -# Protein -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_coo_sendlens.out - -srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_coo_sendlens.out - -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Protein --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_protein_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out - -# # Amazon -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset Amazon --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_amazon_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out -# -# # Reddit -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --timing --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 1 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c1_rowsel_bsend_nobarrier.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 2 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c2_rowsel_bsend_nobarrier.out -# -# srun -l -n 36 --ntasks-per-node 4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset reddit --batch-size 512 --samp-num 10 --n-bulkmb 256 --replication 6 --sample-method sage --baseline --hostname $hostname &> tmp_reddit_sage_bs512sn10_mb256_p36c6_rowsel_bsend_nobarrier.out diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index daaa0c3..901e658 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -123,6 +123,15 @@ at::Tensor sparse_coo_tensor_gpu(const at::Tensor& indices, sparse_dim, dense_dim, size, indices, values, values.options().layout(at::kSparse)); } +at::Tensor sparse_csr_tensor_gpu(const at::Tensor& crows, + const at::Tensor& cols, + const at::Tensor& values, + at::ArrayRef size) { + + // return at::sparse_compressed_tensor(crows, cols, values, size, c10::Layout::SparseCsr); // torch 1.13 + return at::sparse_csr_tensor(crows, cols, values, size, values.options().layout(at::kSparseCsr)); // 1.11 +} + template void printCusparseDnMat(int64_t rows, int64_t cols, int64_t ld, T *values_dev) { T* values_host = new T[rows*cols]; @@ -1123,7 +1132,8 @@ void rowselect_csr_gpu(const at::Tensor& nnz_cols, const at::Tensor& row_offsets } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("sparse_coo_tensor_gpu", &sparse_coo_tensor_gpu, "Sparse Tensor GPU-only constructor"); + m.def("sparse_coo_tensor_gpu", &sparse_coo_tensor_gpu, "Sparse COO Tensor GPU-only constructor"); + m.def("sparse_csr_tensor_gpu", &sparse_csr_tensor_gpu, "Sparse CSR Tensor GPU-only constructor"); m.def("spmm_gpu", &spmm_gpu, "SpMM wrapper for cusparse"); m.def("spgemm_gpu", &spgemm_gpu, "SpGEMM wrapper for cusparse"); m.def("downsample_gpu", &downsample_gpu, "Downsampling for LADIES sampling algorithm"); diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index 7825fde..a765a24 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -246,6 +246,7 @@ def test_nccl(args): else: recv_tensor = torch.cuda.FloatTensor(buffer_size) dist.recv(recv_tensor, src=0) + dist.barrier() dist.barrier() torch.cuda.nvtx.range_pop() @@ -284,35 +285,37 @@ def test_nccl(args): 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 + 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) + recv_tensor = torch.cuda.FloatTensor(buffer_size) - # if rank == 0 and r == run_count - 1: - # start_time(start_timer) + if rank == 0 and r == run_count - 1: + start_time(start_timer) - # dist.scatter(recv_tensor, tensors, src=0) - # torch.cuda.synchronize() + 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() - # torch.cuda.nvtx.range_pop() + 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}") From 7f32c65b858a5dba3e93685214388ac11c745ea5 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 18 Feb 2023 16:37:41 -0800 Subject: [PATCH 119/178] All CUDA devices in node visible to rank (fixes nccl bandwidth issue) --- cagnet/samplers/sage.py | 13 ++++++------- cagnet/samplers/utils.py | 40 ++++++++++++++++++++++++++-------------- examples/gcn_15d.py | 5 +++-- tests/test_nccl.cu | 4 ++-- tests/test_nccl_p2p.py | 2 ++ 5 files changed, 39 insertions(+), 25 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 25a4cdb..6ca79f6 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -67,6 +67,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # adj_matrices[i][j] -- mb i layer j adj_matrices = [None] * n_layers # adj_matrices[i] -- bulk mb mtx for layer j + gpu = torch.device(f"cuda:{torch.cuda.current_device()}") + if not baseline_compare: # start_time(total_start_timer) total_start_timer.record() @@ -79,11 +81,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # Expand batches matrix if baseline_compare: total_start_timer.record() - batches_expand_rows = torch.arange(mb_count * nnz, \ - device=torch.device("cuda:0")) - batches_expand_idxs = torch.stack( - (batches_expand_rows, batches._indices()[1, :]) - ) + batches_expand_rows = torch.arange(mb_count * nnz, device=gpu) + batches_expand_idxs = torch.stack((batches_expand_rows, batches._indices()[1, :])) batches_expand = torch.sparse_coo_tensor( batches_expand_idxs, batches._values(), @@ -120,9 +119,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, frontier_nnz_sizes = torch.clamp(frontier_nnz_sizes, max=frontier_size) next_frontier_rows = torch.repeat_interleave( - torch.arange(batch_size * mb_count, device=torch.device("cuda:0")), + torch.arange(batch_size * mb_count, device=gpu), frontier_size) - nextf_cols_idxs = torch.arange(next_frontier_nnz.size(0), device=torch.device("cuda:0")) + 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 diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index f203737..9e8d079 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -318,9 +318,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ 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] * 3 * recv_proc_count - # send_idx = 0 + 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): @@ -348,11 +348,21 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) + # 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) + + 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, recv_rank, tag=1) + send_ops[send_idx + 2] = \ + dist.P2POp(dist.isend, matb_send_cols, recv_rank, tag=2) + send_ops[send_idx + 3] = \ + dist.P2POp(dist.isend, matb_send_values, recv_rank, tag=3) + send_idx += 4 else: # matb_recv_crows = matb_send_crows.clone() matb_recv_cols = matb_send_cols.clone() @@ -395,10 +405,10 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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: - # reqs = dist.batch_isend_irecv(send_ops) - # for req in reqs: - # req.wait() + if len(send_ops) > 0: + 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}"] = [] @@ -512,6 +522,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk = matc_chunk.coalesce() matc += matc_chunk 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, \ @@ -594,6 +605,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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: @@ -610,8 +622,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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=torch.device("cuda:0")), - matc_chunk_counts) + 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) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 9d2278e..b9c44cf 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -346,6 +346,7 @@ def main(args, batches=None): 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) # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) edge_count = adj_matrix.size(1) @@ -613,8 +614,8 @@ def main(args, batches=None): help="sampling algorithm for training") 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("--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, diff --git a/tests/test_nccl.cu b/tests/test_nccl.cu index 1b05cbb..e7d2977 100644 --- a/tests/test_nccl.cu +++ b/tests/test_nccl.cu @@ -131,11 +131,11 @@ int main(int argc, char* argv[]) cudaEventRecord(start); } - ncclGroupStart(); + // ncclGroupStart(); for (int j = 0; j < nRanks - 1; j++) { NCCLCHECK(ncclSend((const void*)sendbuffs[j], size, ncclFloat, j + 1, comm, s)); } - ncclGroupEnd(); + // ncclGroupEnd(); //communicating using NCCL // NCCLCHECK(ncclAllReduce((const void*)sendbuff, (void*)recvbuff, size, ncclFloat, ncclSum, // comm, s)); diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index a765a24..4d2039a 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -27,6 +27,8 @@ def test_nccl(args): size = dist.get_world_size() print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}") print(f"backend: {dist.get_backend()}") + print(f"rank: {rank} visible_devices: {os.environ['CUDA_VISIBLE_DEVICES']}") + torch.cuda.set_device(rank) start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) From 60e1362ee55c24d1ec96629d1ac5060fdc7268b1 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 20 Feb 2023 18:20:36 -0800 Subject: [PATCH 120/178] CSR allreduce --- cagnet/samplers/sage.py | 1 + cagnet/samplers/utils.py | 262 ++++++++++++++++++++++++++------------- examples/gcn_15d.py | 1 + tests/test_sage.py | 2 + 4 files changed, 181 insertions(+), 85 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 6ca79f6..e1ce827 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -97,6 +97,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() + batches = batches.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 9e8d079..6c03cb7 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -34,6 +34,69 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=Fals 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, timing_dict=None): + while left < right: + 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) + if rank <= mid_rank: + 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() + else: + 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.LongTensor(mat.size(0) + 1) + cols_recv = torch.cuda.LongTensor(nnz_recv.item()) + vals_recv = torch.cuda.FloatTensor(nnz_recv.item()) + + ops = [None] * 6 + if rank <= mid_rank: + ops[0] = dist.P2POp(dist.isend, mat.crow_indices().long(), recv_rank, tag=0) + ops[1] = dist.P2POp(dist.isend, mat.col_indices().long(), 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() + else: + ops[0] = dist.P2POp(dist.isend, mat.crow_indices().long(), recv_rank, tag=0) + ops[1] = dist.P2POp(dist.isend, mat.col_indices().long(), 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() + + torch.cuda.synchronize() + mat_recv = torch.sparse_csr_tensor(rows_recv, cols_recv, vals_recv, size=mat.size()) + mat = mat + mat_recv + + if rank <= mid_rank: + right = mid_rank + else: + left = mid_rank + 1 + return mat + def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, name, timing_dict=None): start_timer = torch.cuda.Event(enable_timing=True) @@ -405,7 +468,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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: + if len(send_ops) > 0 and matb.layout == torch.sparse_csr: reqs = dist.batch_isend_irecv(send_ops) for req in reqs: req.wait() @@ -522,7 +586,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk = matc_chunk.coalesce() matc += matc_chunk stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-loc-spspmm-{name}") - matc = matc.to_sparse_csr() + # 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, \ @@ -532,6 +596,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}"] = [] @@ -548,103 +620,123 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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']}") - torch.cuda.nvtx.range_push("nvtx-matc-coalesce") - 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) - torch.cuda.nvtx.range_pop() # matc-coalesce - # Implement sparse allreduce w/ all_gather and padding - torch.cuda.nvtx.range_push("nvtx-reduce-nnz") - 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_pop() # reduce-nnz + rank_row_start = rank_c * replication + rank_row_stop = (rank_c + 1) * replication - 1 - 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 = matc.to_sparse_csr() + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coo2csr{name}", barrier=True) - 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 + print(f"before csr_allred start: {rank_row_start} {rank_row_stop}", flush=True) + matc = csr_allreduce(matc, rank_row_start, rank_row_stop, rank) + print(f"after csr_allred", flush=True) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) - 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())) + 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() - # 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] + # # Old sparse reduction + # torch.cuda.nvtx.range_push("nvtx-matc-coalesce") + # 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) + # torch.cuda.nvtx.range_pop() # matc-coalesce + + # # Implement sparse allreduce w/ all_gather and padding + # torch.cuda.nvtx.range_push("nvtx-reduce-nnz") + # 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_pop() # reduce-nnz - 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-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-reduce") - start_time(start_timer) - # matc_recv = torch.stack(matc_recv) - # matc = torch.sparse.sum(matc_recv, dim=0) + # 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())) - 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 + # # 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] - 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 + # 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-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 + # 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 - return matc_nnz_indices, matc_nnz_values + # 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, sa_recv_buff, diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index b9c44cf..41a4d46 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -476,6 +476,7 @@ def main(args, batches=None): col_groups, args.timing, args.baseline) print("second run", flush=True) + nnz_row_masks.fill_(False) torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") # torch.manual_seed(rank_col) diff --git a/tests/test_sage.py b/tests/test_sage.py index 778454e..4479d8a 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -96,6 +96,8 @@ def main(args): parser = argparse.ArgumentParser(description='test LADIES sampling code') parser.add_argument("--dataset", type=str, default="Cora", help="dataset to train") + parser.add_argument("--gpu", type=int, default=4, + help="gpus per node") parser.add_argument("--n-layers", type=int, default=1, help="number of hidden gcn layers") parser.add_argument("--batch-size", type=int, default=512, From 9ef45a0d181bd4d517590d3e54bc11ad7dd93331 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 21 Feb 2023 15:14:42 -0800 Subject: [PATCH 121/178] Minor cleanup for csr_allreduce --- cagnet/samplers/sage.py | 1 + cagnet/samplers/utils.py | 69 ++++++++++++++++------------------------ 2 files changed, 29 insertions(+), 41 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index e1ce827..8819463 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -90,6 +90,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand = batches_expand.to_sparse_csr() adj_matrix = adj_matrix.to_sparse_csr() + print(f"adj_matrix.dtype: {adj_matrix.crow_indices().dtype}", flush=True) p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "sage", diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 6c03cb7..af37033 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -47,46 +47,29 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): ops = [None, None] nnz_send = torch.cuda.IntTensor(1).fill_(mat._nnz()) nnz_recv = torch.cuda.IntTensor(1) - if rank <= mid_rank: - 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() - else: - 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() + + 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.LongTensor(mat.size(0) + 1) - cols_recv = torch.cuda.LongTensor(nnz_recv.item()) + rows_recv = torch.cuda.IntTensor(mat.size(0) + 1) + cols_recv = torch.cuda.IntTensor(nnz_recv.item()) vals_recv = torch.cuda.FloatTensor(nnz_recv.item()) ops = [None] * 6 - if rank <= mid_rank: - ops[0] = dist.P2POp(dist.isend, mat.crow_indices().long(), recv_rank, tag=0) - ops[1] = dist.P2POp(dist.isend, mat.col_indices().long(), 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() - else: - ops[0] = dist.P2POp(dist.isend, mat.crow_indices().long(), recv_rank, tag=0) - ops[1] = dist.P2POp(dist.isend, mat.col_indices().long(), 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() - + 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() + torch.cuda.synchronize() mat_recv = torch.sparse_csr_tensor(rows_recv, cols_recv, vals_recv, size=mat.size()) mat = mat + mat_recv @@ -420,9 +403,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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, recv_rank, tag=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, recv_rank, tag=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, recv_rank, tag=3) send_idx += 4 @@ -491,11 +474,15 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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)) + # matb_recv_lengths = torch.cuda.LongTensor(nnz_cols.size(0)) + matb_recv_lengths = torch.cuda.IntTensor(nnz_cols.size(0)) dist.recv(matb_recv_lengths, tag=1, src=q) + matb_recv_lengths = matb_recv_lengths.long() - matb_recv_cols = torch.cuda.LongTensor(selected_rows_count_recv.item()) + # 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) + matb_recv_cols = matb_recv_cols.long() matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) dist.recv(matb_recv_values, tag=3, src=q) @@ -629,9 +616,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coo2csr{name}", barrier=True) start_time(start_timer) - print(f"before csr_allred start: {rank_row_start} {rank_row_stop}", flush=True) + print(f"before csr_allred start: {rank_row_start} stop: {rank_row_stop}", flush=True) matc = csr_allreduce(matc, rank_row_start, rank_row_stop, rank) - print(f"after csr_allred", flush=True) + print(f"after csr_allred start: {rank_row_start} stop: {rank_row_stop}", flush=True) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) start_time(start_timer) From 43c294200038b7a06a51d56166a48f45649fded1 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 24 Feb 2023 18:12:41 -0800 Subject: [PATCH 122/178] Memory optimizations --- cagnet/samplers/sage.py | 1 - cagnet/samplers/utils.py | 2 -- examples/gcn_15d.py | 55 ++++++++++++++++++++-------------------- 3 files changed, 28 insertions(+), 30 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 8819463..e1ce827 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -90,7 +90,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand = batches_expand.to_sparse_csr() adj_matrix = adj_matrix.to_sparse_csr() - print(f"adj_matrix.dtype: {adj_matrix.crow_indices().dtype}", flush=True) p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "sage", diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index af37033..4d509d9 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -616,9 +616,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-coo2csr{name}", barrier=True) start_time(start_timer) - print(f"before csr_allred start: {rank_row_start} stop: {rank_row_stop}", flush=True) matc = csr_allreduce(matc, rank_row_start, rank_row_stop, rank) - print(f"after csr_allred start: {rank_row_start} stop: {rank_row_stop}", flush=True) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) start_time(start_timer) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 41a4d46..f0e2e67 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -174,24 +174,25 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep # 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], 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) + # 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) for i in range(len(am_partitions)): proc_node_count = vtx_indices[i + 1] - vtx_indices[i] @@ -368,13 +369,13 @@ def main(args, batches=None): print("end 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) + # for i in range(len(ampbyp)): + # ampbyp[i] = ampbyp[i].t().coalesce().to(device) print("coalescing", flush=True) - g_loc = g_loc.coalesce().t() + g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) + g_loc = g_loc.to(device) g_loc = row_normalize(g_loc) print("done normalizing", flush=True) torch.set_printoptions(precision=10) @@ -436,7 +437,7 @@ def main(args, batches=None): print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}", flush=True) # do it once before timing - torch.manual_seed(0) + # torch.manual_seed(0) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) @@ -465,7 +466,7 @@ def main(args, batches=None): sa_recv_buff, rank, size, row_groups, col_groups, args.timing) elif args.sample_method == "sage": - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) print("first (warmup) run") current_frontier, next_frontier, adj_matrices = \ sage_sampler(g_loc, batches_loc, args.batch_size, \ @@ -475,7 +476,7 @@ def main(args, batches=None): sa_recv_buff, rank, size, row_groups, col_groups, args.timing, args.baseline) - print("second run", flush=True) + print("second runs", flush=True) nnz_row_masks.fill_(False) torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") @@ -530,8 +531,8 @@ def main(args, batches=None): adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code - return current_frontier, next_frontier, adj_matrices, adj_matrix - # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + # return current_frontier, next_frontier, adj_matrices, adj_matrix + return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups return # create GCN model From 144841943440974e3f3eb5e7ec54f04d6109e453 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 27 Feb 2023 04:25:23 -0800 Subject: [PATCH 123/178] minor ladies tweak --- cagnet/samplers/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 4d509d9..44c53cc 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -752,7 +752,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication torch.cuda.nvtx.range_push("nvtx-compute-p") p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) if name == "ladies": - p_num_values = torch.square(p_num_values) + p_num_values = torch.square(p_num_values).double() elif name == "sage": p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) From 1bb4bf4291ddcbc39ad5db2739a164ffa6de14f4 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 1 Mar 2023 00:48:31 -0800 Subject: [PATCH 124/178] CSR LADIES --- cagnet/samplers/ladies.py | 6 ++++++ cagnet/samplers/utils.py | 33 +++++++++++++++++++++++++++++++-- examples/gcn_15d.py | 15 +++++++++------ tests/test_ladies.py | 4 ++++ 4 files changed, 50 insertions(+), 8 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index f86aa47..61a5e02 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -55,11 +55,17 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota else: nnz = current_frontier[0, :].size(0) + batches = batches.to_sparse_csr() + adj_matrix = adj_matrix.to_sparse_csr() p = gen_prob_dist(batches, adj_matrix, mb_count, node_count_total, replication, rank, size, \ row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies", timing) 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, batches, sa_masks, sa_recv_buff, nnz, \ diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 44c53cc..41a6f3a 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -244,7 +244,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-matc-inst") start_time(start_timer) chunk_size = math.ceil(float(mata.size(1) / (size / replication))) - if True or mata.layout == torch.sparse_coo: + if 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) @@ -487,6 +487,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_values = torch.cuda.DoubleTensor(selected_rows_count_recv.item()) dist.recv(matb_recv_values, tag=3, src=q) + # the spgemm is better with COO for SAGE prob if not (name == "prob" and alg == "sage"): matb_rows_sum = torch.cuda.LongTensor(chunk_col_size).fill_(0) if matb_recv_lengths.size(0) > 0: @@ -548,7 +549,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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: + if mata.layout == torch.sparse_csr and name == "prob" and alg == "sage": 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)) @@ -574,6 +575,34 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matc += matc_chunk 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) + # 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 + print(f"matc.size: {matc.size()} mata_recv.size: {mata_recv.size()} matb_recv.size: {matb_recv.size()}", flush=True) + matc += torch.mm(mata_recv, matb_recv) + 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, \ diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index f0e2e67..cfa8ae1 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -148,11 +148,12 @@ def row_normalize(mx): 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().double(), - mx._indices(), mx._values().double(), + print(f"r_mat_inv.size: {r_mat_inv.size()} mx.size: {mx.size()}", flush=True) + 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, size=(r_mat_inv.size(0), mx.size(1))) + 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, \ @@ -349,7 +350,8 @@ def main(args, batches=None): print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) torch.cuda.set_device(rank % args.gpu) - # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + if args.sample_method == "ladies": + adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) edge_count = adj_matrix.size(1) partitioning = Partitioning.ONE5D @@ -376,6 +378,7 @@ def main(args, batches=None): g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) g_loc = g_loc.to(device) + print(f"g_loc: {g_loc}", flush=True) g_loc = row_normalize(g_loc) print("done normalizing", flush=True) torch.set_printoptions(precision=10) @@ -531,8 +534,8 @@ def main(args, batches=None): adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code - # return current_frontier, next_frontier, adj_matrices, adj_matrix - return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + return current_frontier, next_frontier, adj_matrices, adj_matrix + # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups return # create GCN model diff --git a/tests/test_ladies.py b/tests/test_ladies.py index 9ded29b..d09332a 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -115,6 +115,10 @@ def main(args): 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) From ee1aec93894946a30e380b9f6e2064cfbb33f98d Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 13 Mar 2023 13:29:48 -0700 Subject: [PATCH 125/178] Non SpGEMM col extraction for c=1 --- cagnet/samplers/ladies.py | 2 +- cagnet/samplers/utils.py | 227 +++++++++++++++++++++++--------------- examples/gcn_15d.py | 14 ++- tests/test_ladies.py | 1 + 4 files changed, 151 insertions(+), 93 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 61a5e02..8824acc 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -78,7 +78,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) 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": + 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 = [] diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 41a6f3a..41b98f8 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -80,6 +80,71 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): left = mid_rank + 1 return mat +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 + for i in range(stages): + q = stages * rank_col + i + # rows = torch.arange(q * batch_size, (q + 1) * batch_size).cuda() + # row_mask = torch.cuda.BoolTensor(sample_mtx._nnz()).fill_(0) + # rowselect_csr_gpu(rows, sample_mtx.crow_indices(), row_mask, batch_size, sample_mtx._nnz()) + + # row_lengths = sample_mtx.crow_indices()[1:] - sample_mtx.crow_indices()[:-1] + # row_lengths = row_lengths[rows] + # crow_indices = torch.cuda.LongTensor(batch_size + 1) + # crow_indices[0] = 0 + # crow_indices[1:] = torch.cumsum(row_lengths, dim=0) + 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)])) + batch_mtx = batch_mtx.t() + + 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) + + 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)])) + + 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 = csr_allreduce(sample_mtx, rank_row_start, rank_row_stop, rank) + sample_mtx = sample_mtx.to_sparse_coo() + 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) @@ -93,7 +158,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + 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) @@ -101,7 +166,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-bcast-nnz-{name}", barrier=True) start_time(start_timer) if q == rank: @@ -110,12 +175,12 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + 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}") + 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 @@ -127,63 +192,31 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + 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) - - # print(f"before spgemm_gpu", flush=True) - # print(f"mata_chunk_indices[0,:]: {mata_chunk_indices[0,:].int()}", flush=True) - # print(f"mata_chunk_indices[0,:].size: {mata_chunk_indices[0,:].int().size()}", flush=True) - # print(f"mata.size(0): {mata.size(0)}", flush=True) - # mata_chunk_indices = mata_chunk_indices.int() - # matb_recv_indices = matb_recv_indices.int() - # print(f"matb_recv_indices.row: {matb_recv_indices[0,:]}", flush=True) - # print(f"matb_recv_values: {matb_recv_values}", flush=True) - # matc_outputs = spgemm_gpu(mata_chunk_indices[0,:], # A_rows - # mata_chunk_indices[1,:], # A_cols - # mata_chunk_values.float(), # A_vals - # matb_recv_indices[0,:], # B_rows - # matb_recv_indices[1,:], # B_cols - # matb_recv_values.float(), # B_vals - # mata.size(0), # n - # chunk_col_size, # k - # matb.size(1)) # m - # print(f"after spgemm_gpu", flush=True) - # print(f"matc_outputs: {matc_outputs}", flush=True) - # mata_chunk_indices = mata_chunk_indices.long() - # matb_recv_indices = matb_recv_indices.long() - # matc_chunk_rows = matc_outputs[0].long() - # matc_chunk_cols = matc_outputs[1].long() - # matc_chunk_values = matc_outputs[2].double() - - # print(f"before matc_chunk_rows: {matc_chunk_rows}") - # matc_chunk_counts = torch.diff(matc_chunk_rows) - # matc_chunk_rows = torch.repeat_interleave(torch.arange(0, mata.size(0), device=torch.device("cuda:0")), - # matc_chunk_counts) - # print(f"after matc_chunk_rows: {matc_chunk_rows}") - # matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) - stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}") + 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}") + 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}") + 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}") + 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}") + 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 @@ -191,7 +224,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-nnz-{name}", barrier=True) start_time(start_timer) matc_recv_indices = [] @@ -202,33 +235,33 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, 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}") + 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}") + 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}") + 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}") + 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}") + 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}") + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-unpad-{name}", barrier=True) return matc_nnz_indices, matc_nnz_values @@ -244,7 +277,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.nvtx.range_push("nvtx-matc-inst") start_time(start_timer) chunk_size = math.ceil(float(mata.size(1) / (size / replication))) - if mata.layout == torch.sparse_coo or (name == "prob" and alg == "sage"): + 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) @@ -361,7 +394,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups nnz_cols_count_list = torch.cat(nnz_cols_count_list, dim=0) if matb.layout == torch.sparse_coo: start_time(start_inner_timer) - rowselect_coo_gpu(nnz_col_ids, matb._indices()[0,:], nnz_row_masks, nnz_cols_count_list, \ + 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) @@ -585,22 +618,24 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_time(start_inner_timer) - # mata_recv = mata_recv.to_sparse_coo() - # matb_recv = matb_recv.to_sparse_coo() + 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) - # 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 - print(f"matc.size: {matc.size()} mata_recv.size: {mata_recv.size()} matb_recv.size: {matb_recv.size()}", flush=True) - matc += torch.mm(mata_recv, matb_recv) + 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) + 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: @@ -770,6 +805,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, sa_recv_buff, timing_dict, name) + # p_num_indices, p_num_values = dist_spgemm15D(numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob") torch.cuda.nvtx.range_pop() stop_timer.record() torch.cuda.synchronize() @@ -1108,7 +1144,7 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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() - timing_dict["construct-nextf"].append(stop_time(start_timer, stop_timer)) + 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 @@ -1150,7 +1186,7 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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)) + 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) @@ -1161,31 +1197,35 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ torch.Size([nnz * mb_count, node_count_total])) sa_masks.fill_(0) sa_recv_buff.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, sa_recv_buff, 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)) + 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": - 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) - elif name == "sage": - shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - rank, size, replication, batch_size, node_count_total, mb_count_total, 1) - sampled_indices[1,:] += row_shift - torch.cuda.nvtx.range_pop() - timing_dict["row-select-expand"].append(stop_time(start_timer, stop_timer)) + # # 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": + # 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) + # elif name == "sage": + # shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + # rank, size, replication, batch_size, node_count_total, mb_count_total, 1) + # 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) @@ -1193,8 +1233,10 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 * mb_count_total])) sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, - torch.Size([nnz * mb_count, node_count_total * mb_count_total])) + 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)) @@ -1203,28 +1245,39 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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.size(1)])) + 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) + + # 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) + torch.cuda.nvtx.range_pop() - timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer)) + 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.size(1)])) + 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)) + timing_dict["set-sample"].append(stop_time(start_timer, stop_timer, barrier=True)) torch.cuda.nvtx.range_pop() # nvtx-select-mtxs diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index cfa8ae1..9f4189e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -148,7 +148,6 @@ def row_normalize(mx): r_inv_values, size=(r_inv.size(0), r_inv.size(0))) # mx = r_mat_inv.mm(mx.float()) - print(f"r_mat_inv.size: {r_mat_inv.size()} mx.size: {mx.size()}", flush=True) 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), @@ -378,6 +377,7 @@ def main(args, batches=None): g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) g_loc = g_loc.to(device) + g_loc = g_loc.double() print(f"g_loc: {g_loc}", flush=True) g_loc = row_normalize(g_loc) print("done normalizing", flush=True) @@ -448,7 +448,8 @@ def main(args, batches=None): sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) if args.sample_method == "ladies": - # torch.manual_seed(rank_col) + torch.manual_seed(rank_col) + print("first (warmup) run") current_frontier, next_frontier, adj_matrices = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -457,10 +458,11 @@ def main(args, batches=None): sa_recv_buff, rank, size, row_groups, col_groups, args.timing) - print() + print("second runs", flush=True) + nnz_row_masks.fill_(False) torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - # torch.manual_seed(rank_col) + torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices_bulk = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -529,7 +531,9 @@ def main(args, batches=None): 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() + adj_matrix = adj_matrix.coalesce().cuda().double() + + # Note: Uncomment when testing correctness adj_matrix = row_normalize(adj_matrix) adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code diff --git a/tests/test_ladies.py b/tests/test_ladies.py index d09332a..f98cfc0 100644 --- a/tests/test_ladies.py +++ b/tests/test_ladies.py @@ -58,6 +58,7 @@ def main(args): 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() From ff733c7ae6bfa973ba130f2d573f77c49c9d0381 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 14 Mar 2023 14:53:49 -0700 Subject: [PATCH 126/178] Debugging column extraction --- cagnet/samplers/utils.py | 69 ++++++++++++++++++++++++++-------------- 1 file changed, 46 insertions(+), 23 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 41b98f8..ab7ba3c 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -88,6 +88,14 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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) + + start_time(start_timer) for i in range(stages): q = stages * rank_col + i # rows = torch.arange(q * batch_size, (q + 1) * batch_size).cuda() @@ -99,6 +107,8 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic # crow_indices = torch.cuda.LongTensor(batch_size + 1) # crow_indices[0] = 0 # crow_indices[1:] = torch.cumsum(row_lengths, dim=0) + + 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() @@ -110,8 +120,13 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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) 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()) @@ -121,7 +136,9 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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] @@ -130,7 +147,11 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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 @@ -143,6 +164,8 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic rank_row_stop = (rank_c + 1) * replication - 1 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): @@ -1212,20 +1235,20 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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": - # 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) - # elif name == "sage": - # shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - # rank, size, replication, batch_size, node_count_total, mb_count_total, 1) - # sampled_indices[1,:] += row_shift - # torch.cuda.nvtx.range_pop() - # timing_dict["row-select-expand"].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": + 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) + elif name == "sage": + shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + rank, size, replication, batch_size, node_count_total, mb_count_total, 1) + 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) @@ -1233,10 +1256,10 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 * mb_count_total])) sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, - torch.Size([nnz * mb_count, node_count_total])) + torch.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])) elif name == "sage": # sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, # size=(nnz * mb_count, node_count_total * mb_count_total * nnz)) @@ -1255,13 +1278,13 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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) + 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) + # 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) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer, barrier=True)) From c5c5418434d5aa038b01906ed55cec62ae693727 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 14 Mar 2023 17:25:46 -0700 Subject: [PATCH 127/178] Non-spgemm column extraction --- cagnet/samplers/utils.py | 119 ++++++++++++++++++++++----------------- 1 file changed, 66 insertions(+), 53 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index ab7ba3c..35a698b 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -95,18 +95,12 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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 - # rows = torch.arange(q * batch_size, (q + 1) * batch_size).cuda() - # row_mask = torch.cuda.BoolTensor(sample_mtx._nnz()).fill_(0) - # rowselect_csr_gpu(rows, sample_mtx.crow_indices(), row_mask, batch_size, sample_mtx._nnz()) - - # row_lengths = sample_mtx.crow_indices()[1:] - sample_mtx.crow_indices()[:-1] - # row_lengths = row_lengths[rows] - # crow_indices = torch.cuda.LongTensor(batch_size + 1) - # crow_indices[0] = 0 - # crow_indices[1:] = torch.cumsum(row_lengths, dim=0) start_time(start_inner_timer) next_frontier_unique = next_frontier_select[q].unique() @@ -123,31 +117,50 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic stop_time_add(start_inner_timer, stop_inner_timer, timing_dict, f"spgemm-extract-batch-{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) + 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) - 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) + 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) - 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) + # 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) @@ -1235,20 +1248,20 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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": - 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) - elif name == "sage": - shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - rank, size, replication, batch_size, node_count_total, mb_count_total, 1) - sampled_indices[1,:] += row_shift - torch.cuda.nvtx.range_pop() - timing_dict["row-select-expand"].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": + # 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) + # elif name == "sage": + # shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), + # rank, size, replication, batch_size, node_count_total, mb_count_total, 1) + # 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) @@ -1256,10 +1269,10 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 * mb_count_total])) # sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, - # torch.Size([nnz * mb_count, node_count_total])) + # torch.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])) elif name == "sage": # sample_mtx = torch.sparse_coo_tensor(sampled_indices, sampled_values, # size=(nnz * mb_count, node_count_total * mb_count_total * nnz)) @@ -1278,13 +1291,13 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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) + # 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) + 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) torch.cuda.nvtx.range_pop() timing_dict["col-select-spgemm"].append(stop_time(start_timer, stop_timer, barrier=True)) From 53cb7c00f71350020243d5b0367cde449a91c059 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 14 Mar 2023 18:36:34 -0700 Subject: [PATCH 128/178] Experimenting with dart counts --- cagnet/samplers/utils.py | 11 +++++++---- examples/gcn_15d.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 35a698b..08b637a 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -175,6 +175,7 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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) @@ -885,9 +886,9 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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 // replication + # if rank_col == replication - 1: + # n_darts_col = n_darts - (replication - 1) * n_darts_col n_darts_col = n_darts next_frontier = torch.sparse_coo_tensor(indices=p._indices(), @@ -1276,8 +1277,10 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 * mb_count_total * nnz])) sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, - torch.Size([nnz * mb_count, node_count_total * mb_count_total * nnz])) + 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))) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 9f4189e..a2db66d 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -430,7 +430,7 @@ def main(args, batches=None): 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.n_darts = avg_degree * args.batch_size * (args.replication ** 4) elif args.sample_method == "sage": args.n_darts = avg_degree print(f"n_darts: {args.n_darts}") From d9f2cb5f9f27d0015b1edc5ca7a83a9f4f36a013 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 17 Mar 2023 12:50:27 -0700 Subject: [PATCH 129/178] Matching LADIES accuracy --- cagnet/samplers/ladies.py | 52 ++++++++++++++++++++------------------- cagnet/samplers/utils.py | 4 ++- examples/gcn_15d.py | 9 +++---- 3 files changed, 34 insertions(+), 31 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 8824acc..8d9d3fa 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -56,8 +56,9 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota nnz = current_frontier[0, :].size(0) batches = batches.to_sparse_csr() - adj_matrix = adj_matrix.to_sparse_csr() - p = gen_prob_dist(batches, adj_matrix, mb_count, node_count_total, replication, rank, size, \ + print(f"prob_adj_matrix: {adj_matrix}", flush=True) + adj_matrix_sq = adj_matrix.square().to_sparse_csr() + p = gen_prob_dist(batches, adj_matrix_sq, mb_count, node_count_total, replication, rank, size, \ row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies", timing) next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, \ @@ -77,27 +78,28 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) - 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)}") + 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/utils.py b/cagnet/samplers/utils.py index 08b637a..ba27112 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -854,7 +854,8 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication torch.cuda.nvtx.range_push("nvtx-compute-p") p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) if name == "ladies": - p_num_values = torch.square(p_num_values).double() + # p_num_values = torch.square(p_num_values).double() + p_num_values = p_num_values.double() elif name == "sage": p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) @@ -862,6 +863,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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])) + print(f"p._indices(): {p._indices()[1,:]} p._values: {p._values()}", flush=True) # print(f"p: {p}") print(f"p.nnz: {p._nnz()}", flush=True) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index a2db66d..3115672 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -378,8 +378,7 @@ def main(args, batches=None): print("normalizing", flush=True) g_loc = g_loc.to(device) g_loc = g_loc.double() - print(f"g_loc: {g_loc}", flush=True) - g_loc = row_normalize(g_loc) + # g_loc = row_normalize(g_loc) print("done normalizing", flush=True) torch.set_printoptions(precision=10) @@ -421,7 +420,7 @@ def main(args, batches=None): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) batches_values = torch.cuda.DoubleTensor(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) + # g_loc = torch.pow(g_loc, 2) torch.cuda.nvtx.range_pop() print(f"g_loc: {g_loc}") @@ -448,7 +447,7 @@ def main(args, batches=None): sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) if args.sample_method == "ladies": - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) print("first (warmup) run") current_frontier, next_frontier, adj_matrices = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ @@ -462,7 +461,7 @@ def main(args, batches=None): nnz_row_masks.fill_(False) torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) current_frontier, next_frontier, adj_matrices_bulk = \ ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ From a4fc94dc1ed76cab11cbd61fdfaa5803f8738612 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 23 Mar 2023 04:17:25 -0700 Subject: [PATCH 130/178] New SAGE and LADIES extraction --- .gitignore | 1 + cagnet/samplers/sage.py | 92 +++++--- cagnet/samplers/utils.py | 298 ++++++++++++++++++-------- examples/gcn_15d.py | 36 ++-- sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 10 +- 6 files changed, 294 insertions(+), 145 deletions(-) diff --git a/.gitignore b/.gitignore index a4255fe..d2d2b55 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,7 @@ gcn_distr_3d.py gcn.py reddit.py tests/data +tests/dgl/ tests/examples tests/LADIES/ tests/test_nccl diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index e1ce827..bba254f 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -103,17 +103,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, timing_dict, "sage") - if baseline_compare: - total_stop_timer.record() - torch.cuda.synchronize() - total_time = total_start_timer.elapsed_time(total_stop_timer) - print(f"total_time: {total_time}", flush=True) - - batches_select = None - next_frontier_select = None - adj_matrices = None - break # for comparing with quiver - # add explicit 0's to next_frontier next_frontier_nnz = next_frontier._values().nonzero().squeeze() frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) @@ -134,10 +123,46 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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=(batch_size * mb_count, node_count_total)) + 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) + batch_vals = torch.cuda.LongTensor(batches_select.size()).fill_(1) + next_frontier_select_vals = next_frontier._values().view(mb_count * batch_size, frontier_size) + next_frontier_select_vals = torch.cat((next_frontier_select_vals, batch_vals), dim=1).view(-1) + + batch_rows = torch.arange(mb_count * batch_size).cuda().view(mb_count * batch_size, 1) + next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * batch_size, frontier_size) + next_frontier_select_rows = torch.cat((next_frontier_select_rows, batch_rows), dim=1).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()).cuda() + adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * batch_size) + 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 = [torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, + size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size]))] + + if baseline_compare: + total_stop_timer.record() + torch.cuda.synchronize() + total_time = total_start_timer.elapsed_time(total_stop_timer) + print(f"total_time_baseline: {total_time}", flush=True) + + # batches_select = None + # next_frontier_select = None + # adj_matrices = None + break # for comparing with quiver + + # next_frontier = next_frontier.coalesce() # next_frontier._values().fill_(1) @@ -155,26 +180,27 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - 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": - 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)}") + 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": + 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/utils.py b/cagnet/samplers/utils.py index ba27112..cfea875 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -57,7 +57,10 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): torch.cuda.synchronize() rows_recv = torch.cuda.IntTensor(mat.size(0) + 1) cols_recv = torch.cuda.IntTensor(nnz_recv.item()) - vals_recv = torch.cuda.FloatTensor(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) @@ -80,12 +83,13 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): left = mid_rank + 1 return mat -def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replication, rank, size, \ - row_groups, col_groups, name, timing_dict=None): +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 = mb_count // replication + stages = (bulk_items // semibulk_size) // replication + print(f"stages: {stages} bulk_items: {bulk_items} semibulk_size: {semibulk_size}", flush=True) sample_mtxs = [None] * stages @@ -95,92 +99,188 @@ def col_select15D(sample_mtx, next_frontier_select, mb_count, batch_size, replic 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() + 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) - 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_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()[q * batch_size] - stop_nnz = sample_mtx.crow_indices()[(q + 1) * batch_size] + 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([batch_size, sample_mtx.size(1)])) + 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) - 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) + # 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) - 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) + 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), 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 = 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() + 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 = 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() + 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._indices(), sample_mtx._values() + 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): @@ -863,7 +963,6 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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])) - print(f"p._indices(): {p._indices()[1,:]} p._values: {p._values()}", flush=True) # print(f"p: {p}") print(f"p.nnz: {p._nnz()}", flush=True) normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) @@ -1251,20 +1350,24 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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": - # 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) - # elif name == "sage": - # shift_rowselect_gpu(row_shift, sampled_indices[0,:], sampled_values.size(0), - # rank, size, replication, batch_size, node_count_total, mb_count_total, 1) - # sampled_indices[1,:] += row_shift - # torch.cuda.nvtx.range_pop() - # timing_dict["row-select-expand"].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 // (16 * replication) # 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 + print(f"nnz: {nnz} semibulk_size: {semibulk_size}", flush=True) + 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) @@ -1272,17 +1375,17 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 * mb_count_total])) sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, - torch.Size([nnz * mb_count, node_count_total])) + 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 = 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 * mb_count_total * nnz])) - sample_mtx = sparse_coo_tensor_gpu(sampled_indices, sampled_values, - torch.Size([nnz * mb_count, node_count_total])) + # 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))) @@ -1296,13 +1399,22 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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) + 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) + # 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)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 3115672..7571fd3 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -166,6 +166,7 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep 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 @@ -173,7 +174,7 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep with torch.no_grad(): # Column partitions am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) - + torch.cuda.synchronize() # 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) @@ -349,9 +350,13 @@ def main(args, batches=None): print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}", flush=True) torch.cuda.set_device(rank % args.gpu) - if args.sample_method == "ladies": - adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + # if args.sample_method == "ladies": + # print(f"before adj_matrix.size: {adj_matrix.size()}", flush=True) + # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) + # print(f"after adj_matrix.size: {adj_matrix.size()}", flush=True) edge_count = adj_matrix.size(1) + torch.cuda.synchronize() + print(f"reached here", flush=True) partitioning = Partitioning.ONE5D @@ -378,6 +383,9 @@ def main(args, batches=None): print("normalizing", flush=True) g_loc = g_loc.to(device) g_loc = g_loc.double() + 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) @@ -429,7 +437,7 @@ def main(args, batches=None): 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 ** 4) + args.n_darts = avg_degree * args.batch_size * (args.replication ** 2) elif args.sample_method == "sage": args.n_darts = avg_degree print(f"n_darts: {args.n_darts}") @@ -471,14 +479,14 @@ def main(args, batches=None): col_groups, args.timing) elif args.sample_method == "sage": # torch.manual_seed(rank_col) - print("first (warmup) run") - current_frontier, next_frontier, adj_matrices = \ - sage_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, - sa_recv_buff, rank, size, row_groups, - col_groups, args.timing, args.baseline) + # print("first (warmup) run") + # current_frontier, next_frontier, adj_matrices = \ + # sage_sampler(g_loc, batches_loc, args.batch_size, \ + # args.samp_num, args.n_bulkmb, \ + # args.n_layers, args.n_darts, \ + # args.replication, nnz_row_masks, + # sa_recv_buff, rank, size, row_groups, + # col_groups, args.timing, args.baseline) print("second runs", flush=True) nnz_row_masks.fill_(False) @@ -537,8 +545,8 @@ def main(args, batches=None): adj_matrix = torch.pow(adj_matrix, 2) # return here while testing sampling code - return current_frontier, next_frontier, adj_matrices, adj_matrix - # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + # return current_frontier, next_frontier, adj_matrices, adj_matrix + return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups return # create GCN model diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index 43d73c9..ffaa2b8 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', ext_modules=[cpp_extension.CUDAExtension('sparse_coo_tensor_cpp', ['sparse_coo_tensor.cu'], - extra_compile_args=["-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/include/", "-lcusparse"])], + extra_compile_args=["-O2", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/include/", "-lcusparse"])], cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 901e658..2707fc4 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -910,7 +910,7 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: __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 shift_size, int semibulk_size) { long id = blockIdx.x * blockDim.x + threadIdx.x; long stride = blockDim.x * gridDim.x; @@ -920,13 +920,14 @@ __global__ void ShiftRowSelect(long *row_shift, long *row_select_rows, int rank, 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 * 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 mb_count, int shift_size, int semibulk_size) { int BLOCK_SIZE = 256; @@ -942,7 +943,8 @@ void shift_rowselect_gpu(const at::Tensor& row_shift, const at::Tensor& row_sele batch_size, node_count, mb_count, - shift_size); + shift_size, + semibulk_size); CHECK_ERROR("shift row select error") } From 6612faa3f7e35478b7d0917863ad3d6bb2fb40e8 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 23 May 2023 13:02:10 -0700 Subject: [PATCH 131/178] Removed unncessary format conversions --- cagnet/samplers/ladies.py | 13 +- cagnet/samplers/sage.py | 27 ++-- cagnet/samplers/utils.py | 62 +++++---- examples/gcn_15d.py | 279 +++++++++++++++++++++++++------------- tests/test_nccl_p2p.py | 30 ++++ 5 files changed, 267 insertions(+), 144 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 8d9d3fa..e74dc14 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -24,7 +24,8 @@ def stop_time(start_timer, stop_timer, barrier=False): return time_taken def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ - replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups, timing): + semibulk_size, replication, sa_masks, rank, size, \ + row_groups, col_groups, timing): total_start_timer = torch.cuda.Event(enable_timing=True) total_stop_timer = torch.cuda.Event(enable_timing=True) @@ -56,22 +57,22 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota nnz = current_frontier[0, :].size(0) batches = batches.to_sparse_csr() - print(f"prob_adj_matrix: {adj_matrix}", flush=True) adj_matrix_sq = adj_matrix.square().to_sparse_csr() p = gen_prob_dist(batches, adj_matrix_sq, mb_count, node_count_total, replication, rank, size, \ - row_groups, col_groups, sa_masks, sa_recv_buff, timing_dict, "ladies", timing) + row_groups, col_groups, sa_masks, timing_dict, "ladies", timing) 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() + # 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, batches, sa_masks, sa_recv_buff, nnz, \ + 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") + replication, rank, size, row_groups, col_groups, timing_dict, i, "ladies", + semibulk_size) adj_matrices[i] = adj_matrix_sample current_frontier = next_frontier diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index bba254f..63ce6b7 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -32,7 +32,7 @@ def stop_time(start_timer, stop_timer, barrier=False): return 0.0 def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ - replication, sa_masks, sa_recv_buff, rank, size, row_groups, col_groups, + replication, sa_masks, rank, size, row_groups, col_groups, timing_arg, baseline_arg): global timing @@ -69,6 +69,16 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, gpu = torch.device(f"cuda:{torch.cuda.current_device()}") + batches_expand_rows = torch.arange(mb_count * batch_size, 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 = batches_expand.to_sparse_csr() + adj_matrix = adj_matrix.to_sparse_csr() + if not baseline_compare: # start_time(total_start_timer) total_start_timer.record() @@ -81,18 +91,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # Expand batches matrix if baseline_compare: total_start_timer.record() - batches_expand_rows = torch.arange(mb_count * nnz, 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 * nnz, node_count_total)) - - batches_expand = batches_expand.to_sparse_csr() - adj_matrix = adj_matrix.to_sparse_csr() p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, replication, rank, size, row_groups, col_groups, - sa_masks, sa_recv_buff, timing_dict, "sage", + sa_masks, timing_dict, "sage", timing_arg) adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: @@ -103,6 +104,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, timing_dict, "sage") + start_time(start_timer) # add explicit 0's to next_frontier next_frontier_nnz = next_frontier._values().nonzero().squeeze() frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) @@ -150,6 +152,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrices = [torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size]))] + timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) if baseline_compare: total_stop_timer.record() @@ -167,7 +170,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # next_frontier._values().fill_(1) batches_select, next_frontier_select, adj_matrix_sample = \ - select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, + 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, i, "sage") diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index cfea875..06f45c6 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -25,7 +25,7 @@ def stop_time(start_timer, stop_timer, barrier=False): torch.cuda.synchronize() time_taken = start_timer.elapsed_time(stop_timer) if barrier: - dist.barrier() + dist.barrier(torch.distributed.group.WORLD) return time_taken else: return 0.0 @@ -89,7 +89,6 @@ def col_select15D(sample_mtx, col_select_mtx, bulk_items, node_count_total, bulk rank_c = rank // replication rank_col = rank % replication stages = (bulk_items // semibulk_size) // replication - print(f"stages: {stages} bulk_items: {bulk_items} semibulk_size: {semibulk_size}", flush=True) sample_mtxs = [None] * stages @@ -103,6 +102,7 @@ def col_select15D(sample_mtx, col_select_mtx, bulk_items, node_count_total, bulk 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 @@ -139,10 +139,11 @@ def col_select15D(sample_mtx, col_select_mtx, bulk_items, node_count_total, bulk 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), coalesced=True) + 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() @@ -403,7 +404,7 @@ def dist_spgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, return matc_nnz_indices, matc_nnz_values def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups, \ - name, nnz_row_masks, matb_recv_buff, timing_dict=None, alg=None): + 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) @@ -925,7 +926,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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, sa_recv_buff, + rank, size, row_groups, col_groups, sa_masks, timing_dict, name, timing_arg): global timing @@ -941,7 +942,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication p_num_indices, p_num_values = dist_saspgemm15D( numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, - sa_recv_buff, timing_dict, name) + timing_dict, name) # p_num_indices, p_num_values = dist_spgemm15D(numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob") torch.cuda.nvtx.range_pop() stop_timer.record() @@ -1262,13 +1263,15 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, return next_frontier -def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ +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): + 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": @@ -1334,14 +1337,14 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ row_select_mtx = sparse_coo_tensor_gpu(row_select_mtx_indices, row_select_mtx_values, torch.Size([nnz * mb_count, node_count_total])) sa_masks.fill_(0) - sa_recv_buff.fill_(0) row_select_mtx = row_select_mtx.to_sparse_csr() - adj_matrix = adj_matrix.to_sparse_csr() + # adj_matrix = adj_matrix.to_sparse_csr() + print(f"adj_matrix.layout: {adj_matrix.layout}") sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ row_groups, col_groups, "rowsel", sa_masks, - sa_recv_buff, timing_dict) + timing_dict) row_select_mtx = row_select_mtx.to_sparse_coo() - adj_matrix = adj_matrix.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, @@ -1355,13 +1358,14 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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 // (16 * replication) # number of batches to column extract from in bulk + # 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 - print(f"nnz: {nnz} semibulk_size: {semibulk_size}", flush=True) 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) @@ -1399,22 +1403,22 @@ def select(next_frontier, adj_matrix, batches, sa_masks, sa_recv_buff, nnz, \ 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) + # 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() + 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) + 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)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 7571fd3..5a42f67 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -22,7 +22,7 @@ import cagnet.nn.functional as CAGF import ogb -# from ogb.nodeproppred import PygNodePropPredDataset +from ogb.nodeproppred import PygNodePropPredDataset import socket @@ -255,7 +255,27 @@ def evaluate(model, graph, features, labels, nid, nid_count, ampbyp, group): def main(args, batches=None): # load and preprocess dataset - device = torch.device('cuda') + # 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}') + print(f"device1: {device}") path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', args.dataset) @@ -307,56 +327,54 @@ def main(args, batches=None): 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.startswith("ogb"): - dataset = PygNodePropPredDataset(name=args.dataset, root="../../data") - - 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 = dataset[0] - # data = data.to(device) - # data.x.requires_grad = True - # inputs = data.x.to(device) - inputs = data.x - # 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 - split_idx = dataset.get_idx_split() - train_idx = split_idx['train'].to(device) - - print(f"adj_matrix.size(): {adj_matrix.size()}") - print(f"adj_matrix: {adj_matrix}") - - # 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) + if ("papers100M" in args.dataset and rank % args.gpu == 0) or "products" in args.dataset: + dataset = PygNodePropPredDataset(name=args.dataset, root="../../data") + + 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 = dataset[0] + # data = data.to(device) + # data.x.requires_grad = True + # inputs = data.x.to(device) + inputs = data.x + # 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 + split_idx = dataset.get_idx_split() + train_idx = split_idx['train'].to(device) + + # print(f"adj_matrix.size(): {adj_matrix.size()}") + # print(f"adj_matrix: {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) - edge_count = adj_matrix.size(1) - torch.cuda.synchronize() - print(f"reached here", flush=True) + if args.dataset != "ogbn-papers100M" and args.dataset != "Protein_sg2": + edge_count = adj_matrix.size(1) partitioning = Partitioning.ONE5D @@ -367,25 +385,89 @@ def main(args, batches=None): if rank_c >= (size // args.replication): return + print("start partitioning", flush=True) - features_loc, g_loc, ampbyp = one5d_partition(rank, size, inputs, adj_matrix, data, \ - inputs, num_classes, args.replication, \ - args.normalize) + if args.dataset == "ogbn-papers100M" or args.dataset == "Protein_sg2": + if rank % args.gpu == 0: + # send g_loc.nnz, g_loc, train_idx.len, and train_idx + for i in range(1, args.gpu): + print(f"iteration i: {i}", flush=True) + _, g_loc, _ = one5d_partition(rank + i, size, inputs, adj_matrix, data, inputs, num_classes, \ + args.replication, args.normalize) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t_() + print("normalizing", flush=True) + g_loc = g_loc.to(device) + g_loc = g_loc.double() + edge_count = adj_matrix.size(1) + + dst_gpu = rank + i + + g_loc_meta = torch.cuda.LongTensor(4) + g_loc_meta[0] = g_loc._nnz() + g_loc_meta[1] = g_loc.size(0) + g_loc_meta[2] = g_loc.size(1) + g_loc_meta[3] = adj_matrix.size(1) + dist.send(g_loc_meta, dst=dst_gpu) + dist.send(g_loc._indices(), dst=dst_gpu) + dist.send(g_loc._values(), dst=dst_gpu) + + if args.dataset == "ogbn-papers100M": + train_idx_len = torch.cuda.LongTensor(1).fill_(train_idx.size(0)) + dist.send(train_idx_len, dst=dst_gpu) + print(f"train_idx_len: {train_idx_len} train_idx: {train_idx}", flush=True) + dist.send(train_idx, dst=dst_gpu) + + _, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + args.replication, args.normalize) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t_() + print("normalizing", flush=True) + g_loc = g_loc.to(device) + g_loc = g_loc.double() + edge_count = adj_matrix.size(1) + + del adj_matrix # Comment when testing + del inputs + else: + src_gpu = rank - (rank % args.gpu) + g_loc_meta = torch.cuda.LongTensor(4).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) + + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + size=torch.Size([g_loc_meta[1], g_loc_meta[2]])) + edge_count = g_loc_meta[3].item() + + if args.dataset == "ogbn-papers100M": + train_idx_len = torch.cuda.LongTensor(1).fill_(0) + dist.recv(train_idx_len, src=src_gpu) + + train_idx = torch.cuda.LongTensor(train_idx_len.item()).fill_(0) + dist.recv(train_idx, src=src_gpu) + print(f"recv train_idx: {train_idx}", flush=True) + torch.cuda.synchronize() + + else: + _, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + args.replication, args.normalize) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t_() + print("normalizing", flush=True) + g_loc = g_loc.to(device) + g_loc = g_loc.double() print("end partitioning", flush=True) + print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) - # features_loc = features_loc.to(device) - # for i in range(len(ampbyp)): - # ampbyp[i] = ampbyp[i].t().coalesce().to(device) - - print("coalescing", flush=True) - g_loc = g_loc.coalesce().t_() - print("normalizing", flush=True) - g_loc = g_loc.to(device) - g_loc = g_loc.double() - 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_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) @@ -406,6 +488,7 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() + print(f"device: {device}") if batches is None: print("beginning constructing batches") batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major @@ -422,7 +505,9 @@ def main(args, batches=None): batches_loc = one5d_partition_mb(rank, size, batches, args.replication, 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 = torch.arange(batches_loc.size(0)).to(device) + print(f"device2: {device}") + print(f"batches_indices_rows.device: {batches_indices_rows.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)) @@ -433,17 +518,15 @@ def main(args, batches=None): print(f"g_loc: {g_loc}") node_count = g_loc.size(1) - adj_matrix = adj_matrix.cuda() + # 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 + args.n_darts = avg_degree * 2 print(f"n_darts: {args.n_darts}") - # del adj_matrix # Comment when testing - print(f"rank: {rank} g_loc: {g_loc}") print(f"rank: {rank} batches_loc: {batches_loc}", flush=True) # do it once before timing @@ -452,7 +535,6 @@ def main(args, batches=None): nnz_row_masks.fill_(0) nnz_recv_upperbound = edge_count // (size // args.replication) - sa_recv_buff = torch.cuda.DoubleTensor(3 * nnz_recv_upperbound).fill_(0) if args.sample_method == "ladies": # torch.manual_seed(rank_col) @@ -461,32 +543,32 @@ def main(args, batches=None): ladies_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, - sa_recv_buff, rank, size, row_groups, - col_groups, args.timing) + args.semibulk, args.replication, + nnz_row_masks, rank, size, + row_groups, col_groups, args.timing) print("second runs", flush=True) nnz_row_masks.fill_(False) - torch.cuda.profiler.cudart().cudaProfilerStart() + # torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") # torch.manual_seed(rank_col) 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.replication, nnz_row_masks, - sa_recv_buff, rank, size, row_groups, - col_groups, args.timing) + args.semibulk, args.replication, \ + nnz_row_masks, rank, size, + row_groups, col_groups, args.timing) elif args.sample_method == "sage": - # torch.manual_seed(rank_col) - # print("first (warmup) run") - # current_frontier, next_frontier, adj_matrices = \ - # sage_sampler(g_loc, batches_loc, args.batch_size, \ - # args.samp_num, args.n_bulkmb, \ - # args.n_layers, args.n_darts, \ - # args.replication, nnz_row_masks, - # sa_recv_buff, rank, size, row_groups, - # col_groups, args.timing, args.baseline) + torch.manual_seed(rank_col) + print("first (warmup) run") + current_frontier, next_frontier, adj_matrices = \ + sage_sampler(g_loc, batches_loc, args.batch_size, \ + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, + rank, size, row_groups, + col_groups, args.timing, args.baseline) print("second runs", flush=True) nnz_row_masks.fill_(False) @@ -498,11 +580,11 @@ def main(args, batches=None): args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ args.replication, nnz_row_masks, - sa_recv_buff, rank, size, row_groups, + rank, size, row_groups, col_groups, args.timing, args.baseline) torch.cuda.nvtx.range_pop() - torch.cuda.profiler.cudart().cudaProfilerStop() + # torch.cuda.profiler.cudart().cudaProfilerStop() if current_frontier is not None: # adj_matrices[i][j] -- mb i layer j @@ -532,21 +614,22 @@ def main(args, batches=None): (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) adj_matrices[j][i] = adj_matrix_sample - # print(f"sample: {adj_matrices}") - # 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() + # # print(f"sample: {adj_matrices}") + # # 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() - # Note: Uncomment when testing correctness - adj_matrix = row_normalize(adj_matrix) - adj_matrix = torch.pow(adj_matrix, 2) - # return here while testing sampling code + # # Note: Uncomment when testing correctness + # adj_matrix = row_normalize(adj_matrix) + # adj_matrix = torch.pow(adj_matrix, 2) + # # return here while testing sampling code # return current_frontier, next_frontier, adj_matrices, adj_matrix - return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups + return current_frontier, next_frontier, adj_matrices, col_groups return # create GCN model @@ -668,6 +751,8 @@ def main(args, batches=None): help='number of minibatches to sample 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", diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index 4d2039a..e1fe394 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -355,6 +355,36 @@ def test_nccl(args): torch.cuda.nvtx.range_pop() torch.cuda.profiler.cudart().cudaProfilerStop() + print(f"alltoall", flush=True) + buffer_size = 10 * 2**20 + 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 rank == 0: + start_time(start_timer) + + # dist.scatter(recv_tensor, tensors, src=0) + dist.all_to_all(recv_tensors, tensors) + torch.cuda.synchronize() + + if rank == 0: + 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() + if __name__ == '__main__': parser = argparse.ArgumentParser(description='NCCL P2P test') parser.add_argument('--hostname', default='127.0.0.1', type=str, From a8404b1ea8c9f52501b9ca25c2635ba430ea4463 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 29 May 2023 17:07:23 -0700 Subject: [PATCH 132/178] Refactored frontiers to support multiple layers --- cagnet/samplers/sage.py | 98 ++++++++++++++++++++--------------------- examples/gcn_15d.py | 86 ++++++++++++++++++------------------ tests/test_sage.py | 6 ++- 3 files changed, 97 insertions(+), 93 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 63ce6b7..7b4016a 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -49,8 +49,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, timing_dict = defaultdict(list) - 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) @@ -78,27 +76,35 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand = batches_expand.to_sparse_csr() adj_matrix = adj_matrix.to_sparse_csr() + current_frontier = batches_expand + frontiers = [None] * (n_layers + 1) - if not baseline_compare: - # start_time(total_start_timer) + if baseline_compare: total_start_timer.record() for i in range(n_layers): if i == 0: nnz = batch_size else: - nnz = current_frontier[0, :].size(0) - + current_frontier_nnzmask = current_frontier._values().nonzero().squeeze() + current_frontier_nnzinds = current_frontier._indices()[:, current_frontier_nnzmask] + current_frontier_nnzvals = current_frontier._values()[current_frontier_nnzmask].double() + current_frontier = torch.sparse_coo_tensor(current_frontier_nnzinds, current_frontier_nnzvals, + size=current_frontier.size()) + nnz = current_frontier._nnz() + current_frontier = current_frontier.to_sparse_csr() + + print(f"i: {i} nnz: {nnz}", flush=True) # Expand batches matrix if baseline_compare: total_start_timer.record() - p = gen_prob_dist(batches_expand, adj_matrix, mb_count, node_count_total, + 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) adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() - batches = batches.to_sparse_coo() + # batches = batches.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, @@ -111,7 +117,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, frontier_nnz_sizes = torch.clamp(frontier_nnz_sizes, max=frontier_size) next_frontier_rows = torch.repeat_interleave( - torch.arange(batch_size * mb_count, device=gpu), + torch.arange(nnz * mb_count, device=gpu), + # torch.arange(batch_size * mb_count, device=gpu), frontier_size) nextf_cols_idxs = torch.arange(next_frontier_nnz.size(0), device=gpu) frontier_remainder = frontier_size - frontier_nnz_sizes @@ -128,60 +135,52 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # Construct sampled adj matrix next_frontier = torch.sparse_coo_tensor(next_frontier_idxs, next_frontier_values, - size=(batch_size * mb_count, node_count_total)) - - 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) - batch_vals = torch.cuda.LongTensor(batches_select.size()).fill_(1) - next_frontier_select_vals = next_frontier._values().view(mb_count * batch_size, frontier_size) + size=(nnz * mb_count, node_count_total)) + # size=(batch_size * mb_count, node_count_total)) + + # next_frontier_select = next_frontier._indices()[1,:].view(mb_count * batch_size, frontier_size) + next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) + # batches_select = torch.masked_select(batches._indices()[1,:], \ + # batches._values().bool()).view(mb_count * batch_size, 1) + current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ + current_frontier.values().bool()).view(mb_count * batch_size, 1) + next_frontier_select = torch.cat((next_frontier_select, current_frontier_select), dim=1) + batch_vals = torch.cuda.LongTensor(current_frontier_select.size()).fill_(1) + # next_frontier_select_vals = next_frontier._values().view(mb_count * batch_size, frontier_size) + next_frontier_select_vals = next_frontier._values().view(mb_count * nnz, frontier_size) next_frontier_select_vals = torch.cat((next_frontier_select_vals, batch_vals), dim=1).view(-1) - batch_rows = torch.arange(mb_count * batch_size).cuda().view(mb_count * batch_size, 1) - next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * batch_size, frontier_size) + # batch_rows = torch.arange(mb_count * batch_size).cuda().view(mb_count * batch_size, 1) + batch_rows = torch.arange(mb_count * nnz).cuda().view(mb_count * nnz, 1) + # next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * batch_size, frontier_size) + next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * nnz, frontier_size) next_frontier_select_rows = torch.cat((next_frontier_select_rows, batch_rows), dim=1).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()).cuda() - adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * batch_size) + # adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * batch_size) + adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * nnz) adj_sample_cols = adj_sample_cols[nnz_mask] adj_matrices_indices = torch.stack((adj_sample_rows, adj_sample_cols)) adj_matrices_values = torch.cuda.DoubleTensor(adj_sample_rows.size(0)).fill_(1) - adj_matrices = [torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, - size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size]))] - timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) - - if baseline_compare: - total_stop_timer.record() - torch.cuda.synchronize() - total_time = total_start_timer.elapsed_time(total_stop_timer) - print(f"total_time_baseline: {total_time}", flush=True) - - # batches_select = None - # next_frontier_select = None - # adj_matrices = None - break # for comparing with quiver - - - # next_frontier = next_frontier.coalesce() - # next_frontier._values().fill_(1) - - batches_select, next_frontier_select, adj_matrix_sample = \ - 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, i, "sage") - + # adj_matrices = [torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, + # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size]))] + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, + # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size])) + size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) adj_matrices[i] = adj_matrix_sample + frontiers[i] = current_frontier_select.clone() current_frontier = next_frontier + timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) + frontiers[n_layers] = next_frontier_select.clone() # print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) - total_stop_timer.record() - torch.cuda.synchronize() - total_time = total_start_timer.elapsed_time(total_stop_timer) + if baseline_compare: + total_stop_timer.record() + torch.cuda.synchronize() + total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) if timing: for k, v in sorted(timing_dict.items()): @@ -206,4 +205,5 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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 + # return current_frontier_select, next_frontier_select, adj_matrices + return frontiers, adj_matrices diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 5a42f67..0f8bb1e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -369,6 +369,7 @@ def main(args, batches=None): # print(f"adj_matrix.size(): {adj_matrix.size()}") # print(f"adj_matrix: {adj_matrix}") + adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) # if args.sample_method == "ladies": # print(f"before adj_matrix.size: {adj_matrix.size()}", flush=True) # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) @@ -562,8 +563,7 @@ def main(args, batches=None): elif args.sample_method == "sage": torch.manual_seed(rank_col) print("first (warmup) run") - current_frontier, next_frontier, adj_matrices = \ - sage_sampler(g_loc, batches_loc, args.batch_size, \ + frontiers, adj_matrices = sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ args.replication, nnz_row_masks, @@ -575,61 +575,61 @@ def main(args, batches=None): torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") # torch.manual_seed(rank_col) - current_frontier, next_frontier, adj_matrices_bulk = \ - sage_sampler(g_loc, batches_loc, args.batch_size, \ + frontiers, adj_matrices_bulk = sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ args.n_layers, args.n_darts, \ args.replication, nnz_row_masks, rank, size, row_groups, col_groups, args.timing, args.baseline) - + torch.cuda.nvtx.range_pop() # torch.cuda.profiler.cudart().cudaProfilerStop() - if current_frontier is not None: - # adj_matrices[i][j] -- mb i layer j - rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) - adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] - for i in range(args.n_layers): - for j in range(rank_n_bulkmb): - row_select_min = j * args.batch_size - row_select_max = (j + 1) * args.batch_size - - sampled_indices = adj_matrices_bulk[i]._indices() - sampled_values = adj_matrices_bulk[i]._values() - - sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask] - - if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) - else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) - adj_matrices[j][i] = adj_matrix_sample - - # # print(f"sample: {adj_matrices}") + # adj_matrices[i][j] -- mb i layer j + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] + for i in range(args.n_layers): + for j in range(rank_n_bulkmb): + row_select_min = j * args.batch_size + row_select_max = (j + 1) * args.batch_size + + sampled_indices = adj_matrices_bulk[i]._indices() + sampled_values = adj_matrices_bulk[i]._values() + + sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask] + + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) + adj_matrices[j][i] = adj_matrix_sample + + print(f"sample: {adj_matrices}") # # 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() - - # # Note: Uncomment when testing correctness + 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() + + # Note: Uncomment when testing correctness + # adj_matrix = adj_matrix.to_sparse_coo().double() # adj_matrix = row_normalize(adj_matrix) # adj_matrix = torch.pow(adj_matrix, 2) - # # return here while testing sampling code + # return here while testing sampling code + return frontiers, adj_matrices, adj_matrix # return current_frontier, next_frontier, adj_matrices, adj_matrix # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups - return current_frontier, next_frontier, adj_matrices, col_groups + # return current_frontier, next_frontier, adj_matrices, col_groups return # create GCN model diff --git a/tests/test_sage.py b/tests/test_sage.py index 4479d8a..583fd24 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -88,7 +88,11 @@ def main(args): 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") + 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") From 75c5ac1aecf5c9b4c1c49e700c45ee328f793c6a Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 31 May 2023 16:06:21 -0700 Subject: [PATCH 133/178] Restructures Q matrix per layer SAGE --- cagnet/samplers/sage.py | 32 +++++--- examples/gcn_15d.py | 89 ++++++++++++---------- tests/test_sage.py | 164 +++++++++++++++++++++------------------- 3 files changed, 156 insertions(+), 129 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 7b4016a..14ca549 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -64,6 +64,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # 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()}") @@ -77,19 +78,23 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand = batches_expand.to_sparse_csr() adj_matrix = adj_matrix.to_sparse_csr() current_frontier = batches_expand - frontiers = [None] * (n_layers + 1) if baseline_compare: total_start_timer.record() + for i in range(n_layers): + print(f"Sampling layer {i}") if i == 0: nnz = batch_size else: current_frontier_nnzmask = current_frontier._values().nonzero().squeeze() - current_frontier_nnzinds = current_frontier._indices()[:, current_frontier_nnzmask] + current_frontier_nnzcols = current_frontier._indices()[1, current_frontier_nnzmask] + current_frontier_nnzrows = torch.arange(current_frontier_nnzcols.numel()).cuda() + current_frontier_nnzinds = torch.stack((current_frontier_nnzrows, current_frontier_nnzcols)) current_frontier_nnzvals = current_frontier._values()[current_frontier_nnzmask].double() current_frontier = torch.sparse_coo_tensor(current_frontier_nnzinds, current_frontier_nnzvals, - size=current_frontier.size()) + size=torch.Size([current_frontier_nnzcols.numel(), + current_frontier.size(1)])) nnz = current_frontier._nnz() current_frontier = current_frontier.to_sparse_csr() @@ -97,6 +102,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # Expand batches matrix if baseline_compare: total_start_timer.record() + print(f"current_frontier: {current_frontier}") 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", @@ -143,18 +149,23 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # batches_select = torch.masked_select(batches._indices()[1,:], \ # batches._values().bool()).view(mb_count * batch_size, 1) current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ - current_frontier.values().bool()).view(mb_count * batch_size, 1) - next_frontier_select = torch.cat((next_frontier_select, current_frontier_select), dim=1) + current_frontier.values().bool()).view(current_frontier._nnz(), 1) + if i == 0: + frontiers[0] = current_frontier_select.clone() + + # next_frontier_select = torch.cat((next_frontier_select, current_frontier_select), dim=1) batch_vals = torch.cuda.LongTensor(current_frontier_select.size()).fill_(1) # next_frontier_select_vals = next_frontier._values().view(mb_count * batch_size, frontier_size) - next_frontier_select_vals = next_frontier._values().view(mb_count * nnz, frontier_size) - next_frontier_select_vals = torch.cat((next_frontier_select_vals, batch_vals), dim=1).view(-1) + # next_frontier_select_vals = next_frontier._values().view(mb_count * nnz, frontier_size) + # next_frontier_select_vals = torch.cat((next_frontier_select_vals, batch_vals), dim=1).view(-1) + next_frontier_select_vals = next_frontier._values().view(-1) # batch_rows = torch.arange(mb_count * batch_size).cuda().view(mb_count * batch_size, 1) batch_rows = torch.arange(mb_count * nnz).cuda().view(mb_count * nnz, 1) # next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * batch_size, frontier_size) - next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * nnz, frontier_size) - next_frontier_select_rows = torch.cat((next_frontier_select_rows, batch_rows), dim=1).view(-1) + # next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * nnz, frontier_size) + # next_frontier_select_rows = torch.cat((next_frontier_select_rows, batch_rows), dim=1).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] @@ -171,10 +182,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size])) size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) adj_matrices[i] = adj_matrix_sample - frontiers[i] = current_frontier_select.clone() + frontiers[i + 1] = next_frontier_select.clone() current_frontier = next_frontier timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) - frontiers[n_layers] = next_frontier_select.clone() # print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) if baseline_compare: diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 0f8bb1e..648e548 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -575,7 +575,7 @@ def main(args, batches=None): torch.cuda.profiler.cudart().cudaProfilerStart() torch.cuda.nvtx.range_push("nvtx-sampler") # torch.manual_seed(rank_col) - frontiers, adj_matrices_bulk = sage_sampler(g_loc, batches_loc, args.batch_size, \ + 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, \ args.replication, nnz_row_masks, @@ -585,52 +585,57 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() # torch.cuda.profiler.cudart().cudaProfilerStop() - # adj_matrices[i][j] -- mb i layer j + # adj_matrices[i][j] -- layer i mb j rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) - adj_matrices = [[None] * args.n_layers for x in range(rank_n_bulkmb)] - for i in range(args.n_layers): + adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] + frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] + for i in range(args.n_layers + 1): for j in range(rank_n_bulkmb): row_select_min = j * args.batch_size row_select_max = (j + 1) * args.batch_size - sampled_indices = adj_matrices_bulk[i]._indices() - sampled_values = adj_matrices_bulk[i]._values() - - sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask] - - if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) - else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.batch_size + args.samp_num * args.batch_size)) - adj_matrices[j][i] = adj_matrix_sample - - print(f"sample: {adj_matrices}") - # # 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() - - # Note: Uncomment when testing correctness - # adj_matrix = adj_matrix.to_sparse_coo().double() - # adj_matrix = row_normalize(adj_matrix) - # adj_matrix = torch.pow(adj_matrix, 2) - # return here while testing sampling code - - return frontiers, adj_matrices, adj_matrix - # return current_frontier, next_frontier, adj_matrices, adj_matrix - # return current_frontier, next_frontier, adj_matrices, adj_matrix, col_groups - # return current_frontier, next_frontier, adj_matrices, col_groups - return + if i < args.n_layers: + sampled_indices = adj_matrices_bulk[i]._indices() + sampled_values = adj_matrices_bulk[i]._values() + + sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask] + + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, \ + args.batch_size + args.samp_num * args.batch_size)) + adj_matrices[i][j] = adj_matrix_sample + frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] + frontiers[i][j] = frontiers_sample + + # # 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() + + # Note: Uncomment when testing correctness + # adj_matrix = adj_matrix.to_sparse_coo().double() + # adj_matrix = row_normalize(adj_matrix) + # adj_matrix = torch.pow(adj_matrix, 2) + # return here while testing sampling code + + # print(f"frontiers[0][0].size: {frontiers[0][0].size()}") + # print(f"frontiers[0][1].size: {frontiers[1][0].size()}") + # print(f"frontiers[0][0]: {frontiers[0][0]}") + # print(f"frontiers[0][1]: {frontiers[1][0]}") + # print(f"batches: {batches}") + return frontiers, adj_matrices, adj_matrix # create GCN model torch.manual_seed(0) diff --git a/tests/test_sage.py b/tests/test_sage.py index 583fd24..c7c2c34 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -15,86 +15,98 @@ def main(args): """ # 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}") + frontiers, adj_matrices_all, g_loc = examples.gcn_15d.main(args) + print(f"g_loc: {g_loc}") print("") print("Beginning SAGE sampler test") - 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 {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 - - mb_id = i // args.batch_size - adj_matrices = adj_matrices_all[mb_id] # adj_matrices for minibatch i - - vtx_mask = adj_matrices[0]._indices()[0,:] == (i % args.batch_size) - vtx_indices = adj_matrices[0]._indices()[:, vtx_mask] - vtx_values = adj_matrices[0]._values()[vtx_mask] - adj_matrices = [torch.sparse_coo_tensor(vtx_indices, vtx_values, size=adj_matrices[0].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 + for l in range(args.n_layers): + print(f"Testing layer {l}", flush=True) + + current_frontier_all = frontiers[l] + next_frontier_all = frontiers[l + 1] + torch.set_printoptions(edgeitems=25) + size = dist.get_world_size() + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) * args.batch_size + + for i in range(rank_n_bulkmb): + print(f"Testing batch {i}", flush=True) + mb_id = i // args.batch_size + current_frontier = current_frontier_all[mb_id].view(current_frontier_all[mb_id].numel(), 1) + current_frontier = current_frontier[i % args.batch_size] # current_frontier for minibatch i + next_frontier = next_frontier_all[mb_id] + next_frontier = next_frontier[i % args.batch_size].unique(sorted=True) + + adj_matrices = adj_matrices_all[l][mb_id] # adj_matrices for minibatch i + + + vtx_mask = adj_matrices._indices()[0,:] == (i % args.batch_size) + vtx_indices = adj_matrices._indices()[:, vtx_mask] + vtx_values = adj_matrices._values()[vtx_mask] + adj_matrices = [torch.sparse_coo_tensor(vtx_indices, vtx_values, size=adj_matrices.size())] + + agg_neighbors = [] + for v in current_frontier: + vtx = v.item() + neighbors_mask = g_loc._indices()[0,:] == vtx neighbors = (g_loc._indices()[1,:])[neighbors_mask] - 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") + 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"current_frontier: {current_frontier}") + print(f"next_frontier: {next_frontier}") + print(f"agg_neighbors: {agg_neighbors}") + print(f"neighbors[1256]: {g_loc._indices()[1,(g_loc._indices()[0,:] == 1256)]}") + 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') From 6f7094a34b0cdb60a3904f6a3b261e7d113ce3c5 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 31 May 2023 18:53:36 -0700 Subject: [PATCH 134/178] SAGE multi-layer sampling on mb_count>1 --- cagnet/samplers/sage.py | 24 ++++-------------------- examples/gcn_15d.py | 26 +++++++++++--------------- tests/test_sage.py | 2 +- 3 files changed, 16 insertions(+), 36 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 14ca549..a9b3956 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -87,6 +87,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, if i == 0: nnz = batch_size else: + # Restructure current_frontier to only have 1 nnz/row current_frontier_nnzmask = current_frontier._values().nonzero().squeeze() current_frontier_nnzcols = current_frontier._indices()[1, current_frontier_nnzmask] current_frontier_nnzrows = torch.arange(current_frontier_nnzcols.numel()).cuda() @@ -95,14 +96,13 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, current_frontier = torch.sparse_coo_tensor(current_frontier_nnzinds, current_frontier_nnzvals, size=torch.Size([current_frontier_nnzcols.numel(), current_frontier.size(1)])) - nnz = current_frontier._nnz() + nnz = current_frontier._nnz() // mb_count current_frontier = current_frontier.to_sparse_csr() - print(f"i: {i} nnz: {nnz}", flush=True) + print(f"i: {i} nnz: {nnz}") # Expand batches matrix if baseline_compare: total_start_timer.record() - print(f"current_frontier: {current_frontier}") 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", @@ -110,7 +110,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() - # batches = batches.to_sparse_coo() next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, @@ -124,7 +123,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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(batch_size * mb_count, device=gpu), frontier_size) nextf_cols_idxs = torch.arange(next_frontier_nnz.size(0), device=gpu) frontier_remainder = frontier_size - frontier_nnz_sizes @@ -144,42 +142,28 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, size=(nnz * mb_count, node_count_total)) # size=(batch_size * mb_count, node_count_total)) - # next_frontier_select = next_frontier._indices()[1,:].view(mb_count * batch_size, frontier_size) next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) - # batches_select = torch.masked_select(batches._indices()[1,:], \ - # batches._values().bool()).view(mb_count * batch_size, 1) + print(f"next_frontier_select.size: {next_frontier_select.size()}") current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ current_frontier.values().bool()).view(current_frontier._nnz(), 1) if i == 0: frontiers[0] = current_frontier_select.clone() - # next_frontier_select = torch.cat((next_frontier_select, current_frontier_select), dim=1) batch_vals = torch.cuda.LongTensor(current_frontier_select.size()).fill_(1) - # next_frontier_select_vals = next_frontier._values().view(mb_count * batch_size, frontier_size) - # next_frontier_select_vals = next_frontier._values().view(mb_count * nnz, frontier_size) - # next_frontier_select_vals = torch.cat((next_frontier_select_vals, batch_vals), dim=1).view(-1) next_frontier_select_vals = next_frontier._values().view(-1) - # batch_rows = torch.arange(mb_count * batch_size).cuda().view(mb_count * batch_size, 1) batch_rows = torch.arange(mb_count * nnz).cuda().view(mb_count * nnz, 1) - # next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * batch_size, frontier_size) - # next_frontier_select_rows = next_frontier._indices()[0,:].view(mb_count * nnz, frontier_size) - # next_frontier_select_rows = torch.cat((next_frontier_select_rows, batch_rows), dim=1).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()).cuda() - # adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * batch_size) adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * nnz) adj_sample_cols = adj_sample_cols[nnz_mask] adj_matrices_indices = torch.stack((adj_sample_rows, adj_sample_cols)) adj_matrices_values = torch.cuda.DoubleTensor(adj_sample_rows.size(0)).fill_(1) - # adj_matrices = [torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, - # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size]))] adj_matrix_sample = torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, - # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * batch_size])) size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) adj_matrices[i] = adj_matrix_sample frontiers[i + 1] = next_frontier_select.clone() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 648e548..3bd7499 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -585,14 +585,19 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() # torch.cuda.profiler.cudart().cudaProfilerStop() + print(f"frontiers_bulks: {frontiers_bulk}") # adj_matrices[i][j] -- layer i mb j rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] for i in range(args.n_layers + 1): for j in range(rank_n_bulkmb): - row_select_min = j * args.batch_size - row_select_max = (j + 1) * args.batch_size + if i == 0: + row_select_min = j * args.batch_size + row_select_max = (j + 1) * args.batch_size + else: + row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) + row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) if i < args.n_layers: sampled_indices = adj_matrices_bulk[i]._indices() @@ -611,10 +616,12 @@ def main(args, batches=None): else: adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ adj_matrix_sample_values, \ - (args.batch_size, \ - args.batch_size + args.samp_num * args.batch_size)) + (args.batch_size * (args.samp_num ** i), \ + args.batch_size * (args.samp_num ** (i + 1)))) adj_matrices[i][j] = adj_matrix_sample frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] + print(f"i: {i} j: {j}") + print(f"frontiers_sample: {frontiers_sample}") frontiers[i][j] = frontiers_sample # # adj_matrix = None # Uncomment bottom for testing @@ -624,17 +631,6 @@ def main(args, batches=None): size=(node_count, node_count)) adj_matrix = adj_matrix.coalesce().cuda().double() - # Note: Uncomment when testing correctness - # adj_matrix = adj_matrix.to_sparse_coo().double() - # adj_matrix = row_normalize(adj_matrix) - # adj_matrix = torch.pow(adj_matrix, 2) - # return here while testing sampling code - - # print(f"frontiers[0][0].size: {frontiers[0][0].size()}") - # print(f"frontiers[0][1].size: {frontiers[1][0].size()}") - # print(f"frontiers[0][0]: {frontiers[0][0]}") - # print(f"frontiers[0][1]: {frontiers[1][0]}") - # print(f"batches: {batches}") return frontiers, adj_matrices, adj_matrix # create GCN model diff --git a/tests/test_sage.py b/tests/test_sage.py index c7c2c34..7c1f7ad 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -64,7 +64,7 @@ def main(args): print(f"current_frontier: {current_frontier}") print(f"next_frontier: {next_frontier}") print(f"agg_neighbors: {agg_neighbors}") - print(f"neighbors[1256]: {g_loc._indices()[1,(g_loc._indices()[0,:] == 1256)]}") + print(f"neighbors[2662]: {g_loc._indices()[1,(g_loc._indices()[0,:] == 2662)]}") print(f"vtx {vtx} not in aggregated neighborhood") if vertices_match: From fdc8fde3dfc650e811227f36bad1c54cf009dd30 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 4 Aug 2023 19:20:17 -0700 Subject: [PATCH 135/178] Multi-layer sampling (l=2, cora, reddit) --- cagnet/samplers/sage.py | 32 +++++++++++++++++++++----------- cagnet/samplers/utils.py | 14 +++++++++++--- examples/gcn_15d.py | 21 ++++++++++++--------- tests/test_sage.py | 4 ---- 4 files changed, 44 insertions(+), 27 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index a9b3956..90bb086 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -88,18 +88,23 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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() + print(f"before current_frontier: {current_frontier}") + # 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() + 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()[current_frontier_nnzmask].double() + current_frontier_nnzvals = current_frontier._values().double() current_frontier = torch.sparse_coo_tensor(current_frontier_nnzinds, current_frontier_nnzvals, - size=torch.Size([current_frontier_nnzcols.numel(), + size=torch.Size([current_frontier._nnz(), + # size=torch.Size([current_frontier_nnzcols.numel(), current_frontier.size(1)])) - nnz = current_frontier._nnz() // mb_count + # nnz = current_frontier._nnz() // mb_count + nnz = batch_size * (frontier_size ** i) current_frontier = current_frontier.to_sparse_csr() - print(f"i: {i} nnz: {nnz}") # Expand batches matrix if baseline_compare: total_start_timer.record() @@ -115,8 +120,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, timing_dict, "sage") + print(f"before next_frontier: {next_frontier}") start_time(start_timer) - # add explicit 0's to next_frontier + # add explicit 0's to next_frontier to fix the number of rows as bs * sn^l next_frontier_nnz = next_frontier._values().nonzero().squeeze() frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) @@ -130,6 +136,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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) + print(f"nextf_cols_idxs.size: {nextf_cols_idxs.size()}") + print(f"nextf_cols_idxs: {nextf_cols_idxs}") 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)) @@ -141,11 +149,12 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier_values, size=(nnz * mb_count, node_count_total)) # size=(batch_size * mb_count, node_count_total)) + print(f"after next_frontier: {next_frontier}") next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) - print(f"next_frontier_select.size: {next_frontier_select.size()}") - current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ - current_frontier.values().bool()).view(current_frontier._nnz(), 1) + # current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ + # current_frontier.values().bool()).view(current_frontier._nnz(), 1) + current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) if i == 0: frontiers[0] = current_frontier_select.clone() @@ -165,6 +174,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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])) + print(f"before layer {i} adj_matrix_sample: {adj_matrix_sample}") adj_matrices[i] = adj_matrix_sample frontiers[i + 1] = next_frontier_select.clone() current_frontier = next_frontier diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 06f45c6..766baa6 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -943,6 +943,8 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, timing_dict, name) + print(f"p_num_indices: {p_num_indices}") + print(f"p_num_values: {p_num_values}") # p_num_indices, p_num_values = dist_spgemm15D(numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob") torch.cuda.nvtx.range_pop() stop_timer.record() @@ -958,15 +960,18 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # p_num_values = torch.square(p_num_values).double() p_num_values = p_num_values.double() elif name == "sage": - p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) + # p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) + p_num_values = p_num_values.double() scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) # 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()) + normalize_gpu(p_num_values, p_den, p_num_indices[0, :], p_num_values.size(0)) + p_num_values = torch.nan_to_num(p_num_values) p = sparse_coo_tensor_gpu(p_num_indices, p_num_values, torch.Size([numerator.size(0), node_count_total])) - # print(f"p: {p}") print(f"p.nnz: {p._nnz()}", flush=True) - normalize_gpu(p._values(), p_den, p._indices()[0, :], p._nnz()) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) return p @@ -1008,6 +1013,9 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, zero_count = (next_frontier._indices()[0,:] == 0).nonzero().size(0) frontier_nnz_sizes[0] = zero_count + zero_idxs = p._indices()[0, (p._values() == 0)].unique() + frontier_nnz_sizes[zero_idxs] = 0 + torch.cuda.nvtx.range_pop() timing_dict["pre-loop"].append(stop_time(start_timer, stop_timer)) sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 3bd7499..12fede6 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -369,7 +369,7 @@ def main(args, batches=None): # 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, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) # if args.sample_method == "ladies": # print(f"before adj_matrix.size: {adj_matrix.size()}", flush=True) # adj_matrix, _ = add_remaining_self_loops(adj_matrix, num_nodes=inputs.size(0)) @@ -561,7 +561,7 @@ def main(args, batches=None): nnz_row_masks, rank, size, row_groups, col_groups, args.timing) elif args.sample_method == "sage": - torch.manual_seed(rank_col) + # torch.manual_seed(rank_col) print("first (warmup) run") frontiers, adj_matrices = sage_sampler(g_loc, batches_loc, args.batch_size, \ args.samp_num, args.n_bulkmb, \ @@ -585,7 +585,6 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() # torch.cuda.profiler.cudart().cudaProfilerStop() - print(f"frontiers_bulks: {frontiers_bulk}") # adj_matrices[i][j] -- layer i mb j rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] @@ -598,17 +597,21 @@ def main(args, batches=None): else: row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) + adj_row_select_min = j * args.batch_size * (args.samp_num ** i) + adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) if i < args.n_layers: sampled_indices = adj_matrices_bulk[i]._indices() sampled_values = adj_matrices_bulk[i]._values() - sample_select_mask = (row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < row_select_max) + sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < adj_row_select_max) adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - adj_matrix_sample_indices[0,:] -= row_select_min + # adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_indices[0,:] -= adj_row_select_min adj_matrix_sample_values = sampled_values[sample_select_mask] + print(F"i: {i} j: {j} adj_matrix_sample_indices: {adj_matrix_sample_indices}") if args.sample_method == "ladies": adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ adj_matrix_sample_values, \ @@ -618,10 +621,9 @@ def main(args, batches=None): adj_matrix_sample_values, \ (args.batch_size * (args.samp_num ** i), \ args.batch_size * (args.samp_num ** (i + 1)))) + print(f"after layer {i} adj_matrix_sample: {adj_matrix_sample}") adj_matrices[i][j] = adj_matrix_sample frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] - print(f"i: {i} j: {j}") - print(f"frontiers_sample: {frontiers_sample}") frontiers[i][j] = frontiers_sample # # adj_matrix = None # Uncomment bottom for testing @@ -631,7 +633,8 @@ def main(args, batches=None): size=(node_count, node_count)) adj_matrix = adj_matrix.coalesce().cuda().double() - return frontiers, adj_matrices, adj_matrix + # return frontiers, adj_matrices, adj_matrix + return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model torch.manual_seed(0) diff --git a/tests/test_sage.py b/tests/test_sage.py index 7c1f7ad..4fb07ea 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -61,10 +61,6 @@ def main(args): vtx = v.item() if vtx not in agg_neighbors and vtx not in current_frontier: vertices_match = False - print(f"current_frontier: {current_frontier}") - print(f"next_frontier: {next_frontier}") - print(f"agg_neighbors: {agg_neighbors}") - print(f"neighbors[2662]: {g_loc._indices()[1,(g_loc._indices()[0,:] == 2662)]}") print(f"vtx {vtx} not in aggregated neighborhood") if vertices_match: From 9342a4b0715544d2056cc5b48ec95d5895c4876a Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 21 Aug 2023 16:37:28 -0700 Subject: [PATCH 136/178] Mean aggregation support --- cagnet/nn/conv/gcn_conv.py | 62 +++- cagnet/nn/functional.py | 5 +- cagnet/partitionings/partitionings.py | 1 + examples/gcn_15d.py | 501 +++++++++++++------------- sparse-extension/sparse_coo_tensor.cu | 22 ++ 5 files changed, 347 insertions(+), 244 deletions(-) diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index f257592..2472cd9 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -4,7 +4,7 @@ import torch.nn as nn from cagnet.partitionings import Partitioning -from sparse_coo_tensor_cpp import sparse_coo_tensor_gpu, spmm_gpu +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)) @@ -551,7 +551,7 @@ def split3dspmm_dense(self, adj_matrix, inputs, height, middim, width): return z_loc class GCNConv(nn.Module): - def __init__(self, in_feats, out_feats, partitioning, device): + 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) @@ -559,6 +559,7 @@ def __init__(self, in_feats, out_feats, partitioning, device): weight_nonleaf.retain_grad() self.weight = nn.Parameter(weight_nonleaf) self.partitioning = partitioning + self.aggr = aggr def forward(self, gcn, graph, inputs, ampbyp=None): if self.partitioning == Partitioning.ONED: @@ -569,8 +570,65 @@ def forward(self, gcn, graph, inputs, ampbyp=None): 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) +class GCNFuncNONE(torch.autograd.Function): + @staticmethod + def forward(ctx, self, graph, inputs, weight): + # inputs: H + # graph: A + # weight: W + + 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 + + return z + + @staticmethod + def backward(ctx, grad_output): + graph = ctx.graph + inputs, weight = ctx.saved_tensors + self = ctx.self + + ag = torch.cuda.FloatTensor(graph.size(1), grad_output.size(1), \ + device=self.device).fill_(0) + + # Use A^T instead of A implicitly + spmm_gpu(graph.indices()[1].int(), graph.indices()[0].int(), + graph.values(), graph.size(1), + graph.size(0), grad_output, ag) + + if self.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)) + + grad_input = ag.mm(weight.t()) + grad_weight = outer_product(inputs.t(), ag, self.group) + + return None, None, grad_input, grad_weight + class GCNFuncONED(torch.autograd.Function): @staticmethod def forward(ctx, self, graph, ampbyp, inputs, weight): diff --git a/cagnet/nn/functional.py b/cagnet/nn/functional.py index 0d9a139..a726742 100644 --- a/cagnet/nn/functional.py +++ b/cagnet/nn/functional.py @@ -42,7 +42,10 @@ def relu(x, partitioning): return F.relu(x) def log_softmax(self, x, partitioning, dim=1): - if partitioning == Partitioning.ONED or partitioning == Partitioning.ONE5D: + 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) diff --git a/cagnet/partitionings/partitionings.py b/cagnet/partitionings/partitionings.py index 1dbd3ab..fe40743 100644 --- a/cagnet/partitionings/partitionings.py +++ b/cagnet/partitionings/partitionings.py @@ -5,3 +5,4 @@ class Partitioning(enum.IntEnum): ONE5D = 1 TWOD = 2 THREED = 3 + NONE = 4 diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 12fede6..1d6e2fe 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -1,4 +1,5 @@ import argparse +import copy import math import os import os.path as osp @@ -10,9 +11,10 @@ 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 -import torch_geometric.transforms as T +from torch_geometric.loader import NeighborSampler from torch_geometric.utils import add_remaining_self_loops import torch_sparse @@ -20,6 +22,7 @@ from cagnet.partitionings import Partitioning from cagnet.samplers import ladies_sampler, sage_sampler import cagnet.nn.functional as CAGF +import torch.nn.functional as F import ogb from ogb.nodeproppred import PygNodePropPredDataset @@ -27,37 +30,80 @@ import socket class GCN(nn.Module): - def __init__(self, in_feats, n_hidden, n_classes, n_layers, rank, size, partitioning, replication, 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.row_groups = row_groups - self.col_groups = col_groups - self.device = device - self.partitioning = partitioning - self.replication = replication - 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 + 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() + + # 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): + h = inputs + for l, layer in enumerate(self.layers): + print(f"l: {l} graphs[l].size: {graphs[l].size()} h.size: {h.size()}") + h = layer(self, graphs[l], 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 + + @torch.no_grad() + def evaluate(self, graph, features): + subgraph_loader = NeighborSampler(graph._indices(), node_idx=None, + sizes=[-1], batch_size=2048, + 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) + 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_batch = features[n_id].to(self.device) + + h = layer(self, adj_batch, h_batch) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + hs.append(h) + features = torch.cat(hs, dim=0) + return features + + # logits = model(graph, features) + # model.timings = non_eval_timings # don't include evaluation timings + + # 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 get_proc_groups(rank, size, replication): rank_c = rank // replication @@ -218,41 +264,6 @@ def one5d_partition_mb(rank, size, batches, replication, mb_count): batch_partitions = torch.split(batches, math.ceil(float(mb_count / (size / replication))), dim=0) return batch_partitions[rank_c] -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, batches=None): # load and preprocess dataset # Initialize distributed environment with SLURM @@ -275,7 +286,7 @@ def main(args, batches=None): torch.cuda.set_device(rank % args.gpu) device = torch.device(f'cuda:{rank % args.gpu}') - print(f"device1: {device}") + print(f"device1: {device}", flush=True) path = osp.join(osp.dirname(osp.realpath(__file__)), '../..', 'data', args.dataset) @@ -283,10 +294,12 @@ def main(args, batches=None): if args.dataset == "cora": dataset = Planetoid(path, args.dataset, transform=T.NormalizeFeatures()) elif args.dataset == "reddit": - dataset = Reddit(path, transform=T.NormalizeFeatures()) + dataset = Reddit(path) data = dataset[0] + print(f"before data loading", flush=True) data = data.to(device) + print(f"after1 data loading", flush=True) data.x.requires_grad = True inputs = data.x.to(device) inputs.requires_grad = True @@ -379,7 +392,9 @@ def main(args, batches=None): partitioning = Partitioning.ONE5D + print(f"before get_proc_groups", flush=True) row_groups, col_groups = get_proc_groups(rank, size, args.replication) + print(f"after get_proc_groups", flush=True) rank_c = rank // args.replication rank_col = rank % args.replication @@ -453,7 +468,7 @@ def main(args, batches=None): dist.recv(train_idx, src=src_gpu) print(f"recv train_idx: {train_idx}", flush=True) torch.cuda.synchronize() - + features_loc = inputs else: _, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ args.replication, args.normalize) @@ -462,6 +477,7 @@ def main(args, batches=None): print("normalizing", flush=True) g_loc = g_loc.to(device) g_loc = g_loc.double() + features_loc = inputs print("end partitioning", flush=True) print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) @@ -476,9 +492,11 @@ def main(args, batches=None): n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) # rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] + rank_val_mask = torch.split(data.val_mask, n_per_proc, dim=0)[rank_c] # rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] - # labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] + labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] # rank_train_nids = rank_train_mask.nonzero().squeeze() + rank_val_nids = rank_val_mask.nonzero().squeeze() # rank_test_nids = rank_test_mask.nonzero().squeeze() if not args.dataset.startswith("ogb"): @@ -489,152 +507,15 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() - print(f"device: {device}") - if batches is None: - print("beginning constructing batches") - batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major - torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - torch.manual_seed(0) - 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() - print("done constructing batches") - - batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) - - torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") - batches_indices_rows = torch.arange(batches_loc.size(0)).to(device) - print(f"device2: {device}") - print(f"batches_indices_rows.device: {batches_indices_rows.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.DoubleTensor(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) - torch.cuda.nvtx.range_pop() - print(f"g_loc: {g_loc}") - - node_count = g_loc.size(1) + # # # adj_matrix = None # Uncomment bottom for testing # 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 * 2 - print(f"n_darts: {args.n_darts}") - - print(f"rank: {rank} g_loc: {g_loc}") - print(f"rank: {rank} batches_loc: {batches_loc}", flush=True) - # do it once before timing - # torch.manual_seed(0) - nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm - nnz_row_masks.fill_(0) - - nnz_recv_upperbound = edge_count // (size // args.replication) - - if args.sample_method == "ladies": - # torch.manual_seed(rank_col) - print("first (warmup) run") - current_frontier, next_frontier, adj_matrices = \ - 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) - - print("second runs", flush=True) - nnz_row_masks.fill_(False) - # torch.cuda.profiler.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("nvtx-sampler") - # torch.manual_seed(rank_col) - 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) - elif args.sample_method == "sage": - # torch.manual_seed(rank_col) - print("first (warmup) run") - frontiers, adj_matrices = sage_sampler(g_loc, batches_loc, args.batch_size, \ - args.samp_num, args.n_bulkmb, \ - args.n_layers, args.n_darts, \ - args.replication, nnz_row_masks, - rank, size, row_groups, - col_groups, args.timing, args.baseline) - - print("second runs", flush=True) - nnz_row_masks.fill_(False) - torch.cuda.profiler.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("nvtx-sampler") - # torch.manual_seed(rank_col) - 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, \ - args.replication, nnz_row_masks, - rank, size, row_groups, - col_groups, args.timing, args.baseline) - - torch.cuda.nvtx.range_pop() - # torch.cuda.profiler.cudart().cudaProfilerStop() - - # adj_matrices[i][j] -- layer i mb j - rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) - adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] - frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] - for i in range(args.n_layers + 1): - for j in range(rank_n_bulkmb): - if i == 0: - row_select_min = j * args.batch_size - row_select_max = (j + 1) * args.batch_size - else: - row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) - row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) - adj_row_select_min = j * args.batch_size * (args.samp_num ** i) - adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) - - if i < args.n_layers: - sampled_indices = adj_matrices_bulk[i]._indices() - sampled_values = adj_matrices_bulk[i]._values() - - sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < adj_row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - # adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_indices[0,:] -= adj_row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask] - - print(F"i: {i} j: {j} adj_matrix_sample_indices: {adj_matrix_sample_indices}") - if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) - else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size * (args.samp_num ** i), \ - args.batch_size * (args.samp_num ** (i + 1)))) - print(f"after layer {i} adj_matrix_sample: {adj_matrix_sample}") - adj_matrices[i][j] = adj_matrix_sample - frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] - frontiers[i][j] = frontiers_sample - - # # 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() + # adj_matrix = torch.sparse_coo_tensor(adj_matrix, + # torch.cuda.FloatTensor(adj_matrix.size(1)).fill_(1.0), + # size=(node_count, node_count)) + # adj_matrix = adj_matrix.coalesce().cuda().double() # return frontiers, adj_matrices, adj_matrix - return frontiers, adj_matrices, adj_matrix, col_groups + # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model torch.manual_seed(0) @@ -642,18 +523,24 @@ def main(args, batches=None): args.n_hidden, num_classes, args.n_layers, + "mean", rank, size, - partitioning, + Partitioning.NONE, args.replication, device, row_groups=row_groups, col_groups=col_groups) # use optimizer - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) - torch.manual_seed(0) + for W in model.parameters(): + print(W.size()) + print(W) + + dur = [] + # torch.manual_seed(0) total_start = time.time() for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) @@ -661,30 +548,163 @@ def main(args, batches=None): 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_c, col_groups[0], \ - size // args.replication) - - 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) - """ + print("Constructing batches", flush=True) + if batches is None: + batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major + torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") + # torch.manual_seed(0) + 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 = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) + + torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") + batches_indices_rows = torch.arange(batches_loc.size(0)).to(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.DoubleTensor(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) + torch.cuda.nvtx.range_pop() + + 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 * 2 + + print(f"rank: {rank} batches_loc: {batches_loc}", flush=True) + nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm + nnz_row_masks.fill_(0) + + nnz_recv_upperbound = edge_count // (size // args.replication) + + print("Sampling", flush=True) + if args.sample_method == "ladies": + # torch.manual_seed(rank_col) + print("first (warmup) run") + current_frontier, next_frontier, adj_matrices = \ + 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) + + print("second runs", flush=True) + nnz_row_masks.fill_(False) + # torch.cuda.profiler.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("nvtx-sampler") + # torch.manual_seed(rank_col) + 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) + elif args.sample_method == "sage": + nnz_row_masks.fill_(False) + torch.cuda.profiler.cudart().cudaProfilerStart() + torch.cuda.nvtx.range_push("nvtx-sampler") + 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, \ + args.replication, nnz_row_masks, + rank, size, row_groups, + col_groups, args.timing, args.baseline) + + torch.cuda.nvtx.range_pop() + + print("Extracting batches", flush=True) + # adj_matrices[i][j] -- layer i mb j + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] + frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] + for i in range(args.n_layers + 1): + for j in range(rank_n_bulkmb): + if i == 0: + row_select_min = j * args.batch_size + row_select_max = (j + 1) * args.batch_size + else: + row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) + row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) + adj_row_select_min = j * args.batch_size * (args.samp_num ** i) + adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) + + if i < args.n_layers: + sampled_indices = adj_matrices_bulk[i]._indices() + sampled_values = adj_matrices_bulk[i]._values() + + sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < adj_row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + # adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_indices[0,:] -= adj_row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask].float() + + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size * (args.samp_num ** i), \ + args.batch_size * (args.samp_num ** (i + 1)))) + adj_matrices[i][j] = adj_matrix_sample.coalesce() + frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] + frontiers[i][j] = frontiers_sample + + # return frontiers, adj_matrices, adj_matrix, col_groups + g_loc = g_loc.coalesce() + print("Training", flush=True) + for i in range(args.n_bulkmb): + # forward + batch_vtxs = frontiers[0][i].view(-1) + src_vtxs = frontiers[-1][i].view(-1) + features_batch = features_loc[src_vtxs] + adjs = [adj[i].coalesce() for adj in adj_matrices] + adjs.reverse() + logits = model(adjs, features_batch) + print(f"logits: {logits}") + # loss = CAGF.cross_entropy(logits[train_nid], data.y[train_nid], train_nid.size(0), \ + # num_classes, partitioning, rank_c, col_groups[0], \ + # size // args.replication) + loss = F.nll_loss(logits, data.y[batch_vtxs]) + + 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]) + + if rank == 0 and epoch >= 1 and i == args.n_bulkmb - 1: + dur.append(time.time() - total_start) + print("Evaluating...") + # acc = model.evaluate(g_loc, features_loc, labels_rank, rank_val_nids, \ + # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) + out = model.evaluate(g_loc, features_loc) + res = out.argmax(dim=-1) == data.y + acc1 = int(res[data.train_mask].sum()) / int(data.train_mask.sum()) + acc2 = int(res[data.val_mask].sum()) / int(data.val_mask.sum()) + acc3 = int(res[data.test_mask].sum()) / int(data.test_mask.sum()) + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.mean(dur), loss.item(), acc3), 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() @@ -693,7 +713,7 @@ def main(args, batches=None): 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, \ + acc = model.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)) @@ -705,7 +725,6 @@ def main(args, batches=None): # 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}") @@ -723,7 +742,7 @@ def main(args, batches=None): 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, + 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") diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 2707fc4..0fdeb60 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -908,6 +908,27 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: CHECK_ERROR("normalize error") } +__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) { @@ -1153,4 +1174,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From c1962769a176b632cc7bf93a99c370dc96fb64a2 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 22 Aug 2023 12:43:22 -0700 Subject: [PATCH 137/178] Cleanup --- cagnet/samplers/sage.py | 9 +-------- cagnet/samplers/utils.py | 3 --- examples/gcn_15d.py | 30 ++++++++++++++---------------- 3 files changed, 15 insertions(+), 27 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 90bb086..ddd9c51 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -83,12 +83,10 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, total_start_timer.record() for i in range(n_layers): - print(f"Sampling layer {i}") if i == 0: nnz = batch_size else: # Restructure current_frontier to only have 1 nnz/row - print(f"before current_frontier: {current_frontier}") # 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() @@ -120,7 +118,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, timing_dict, "sage") - print(f"before next_frontier: {next_frontier}") 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() @@ -136,8 +133,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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) - print(f"nextf_cols_idxs.size: {nextf_cols_idxs.size()}") - print(f"nextf_cols_idxs: {nextf_cols_idxs}") 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)) @@ -149,7 +144,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier_values, size=(nnz * mb_count, node_count_total)) # size=(batch_size * mb_count, node_count_total)) - print(f"after next_frontier: {next_frontier}") next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) # current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ @@ -174,7 +168,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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])) - print(f"before layer {i} adj_matrix_sample: {adj_matrix_sample}") adj_matrices[i] = adj_matrix_sample frontiers[i + 1] = next_frontier_select.clone() current_frontier = next_frontier @@ -185,7 +178,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, total_stop_timer.record() torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) - print(f"total_time: {total_time}", flush=True) + print(f"total_time: {total_time}", 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": diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 766baa6..b834618 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -943,8 +943,6 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob", sa_masks, timing_dict, name) - print(f"p_num_indices: {p_num_indices}") - print(f"p_num_values: {p_num_values}") # p_num_indices, p_num_values = dist_spgemm15D(numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob") torch.cuda.nvtx.range_pop() stop_timer.record() @@ -971,7 +969,6 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication normalize_gpu(p_num_values, p_den, p_num_indices[0, :], p_num_values.size(0)) p_num_values = torch.nan_to_num(p_num_values) p = sparse_coo_tensor_gpu(p_num_indices, p_num_values, torch.Size([numerator.size(0), node_count_total])) - print(f"p.nnz: {p._nnz()}", flush=True) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) return p diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 1d6e2fe..f82485b 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -58,7 +58,6 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa def forward(self, graphs, inputs): h = inputs for l, layer in enumerate(self.layers): - print(f"l: {l} graphs[l].size: {graphs[l].size()} h.size: {h.size()}") h = layer(self, graphs[l], h) if l != len(self.layers) - 1: h = CAGF.relu(h, self.partitioning) @@ -544,8 +543,8 @@ def main(args, batches=None): total_start = time.time() for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) - if epoch == 1: - total_start = time.time() + if epoch >= 1: + epoch_start = time.time() model.train() print("Constructing batches", flush=True) @@ -581,7 +580,6 @@ def main(args, batches=None): elif args.sample_method == "sage": args.n_darts = avg_degree * 2 - print(f"rank: {rank} batches_loc: {batches_loc}", flush=True) nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm nnz_row_masks.fill_(0) @@ -675,7 +673,6 @@ def main(args, batches=None): adjs = [adj[i].coalesce() for adj in adj_matrices] adjs.reverse() logits = model(adjs, features_batch) - print(f"logits: {logits}") # loss = CAGF.cross_entropy(logits[train_nid], data.y[train_nid], train_nid.size(0), \ # num_classes, partitioning, rank_c, col_groups[0], \ # size // args.replication) @@ -691,17 +688,18 @@ def main(args, batches=None): # 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]) - if rank == 0 and epoch >= 1 and i == args.n_bulkmb - 1: - dur.append(time.time() - total_start) - print("Evaluating...") - # acc = model.evaluate(g_loc, features_loc, labels_rank, rank_val_nids, \ - # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) - out = model.evaluate(g_loc, features_loc) - res = out.argmax(dim=-1) == data.y - acc1 = int(res[data.train_mask].sum()) / int(data.train_mask.sum()) - acc2 = int(res[data.val_mask].sum()) / int(data.val_mask.sum()) - acc3 = int(res[data.test_mask].sum()) / int(data.test_mask.sum()) - print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.mean(dur), loss.item(), acc3), flush=True) + if epoch >= 1: + dur.append(time.time() - epoch_start) + if rank == 0 and (epoch % 5 == 0 or epoch == args.n_epochs - 1): + print("Evaluating", flush=True) + # acc = model.evaluate(g_loc, features_loc, labels_rank, rank_val_nids, \ + # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) + out = model.evaluate(g_loc, features_loc) + res = out.argmax(dim=-1) == data.y + acc1 = int(res[data.train_mask].sum()) / int(data.train_mask.sum()) + acc2 = int(res[data.val_mask].sum()) / int(data.val_mask.sum()) + acc3 = int(res[data.test_mask].sum()) / int(data.test_mask.sum()) + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) total_stop = time.time() print(f"total_time: {total_stop - total_start}") From 8e1f9438f7789329f1b9fa14023cc398cf8c2461 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 6 Sep 2023 16:30:31 -0700 Subject: [PATCH 138/178] Full pipeline for P=1 cora --- cagnet/nn/conv/gcn_conv.py | 66 +++- cagnet/samplers/sage.py | 6 +- cagnet/samplers/utils.py | 67 +--- examples/gcn_15d.py | 527 ++++++++++++++++---------- sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 68 +++- 6 files changed, 454 insertions(+), 282 deletions(-) diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index 2472cd9..3e1347f 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -2,8 +2,10 @@ 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): @@ -561,7 +563,7 @@ def __init__(self, in_feats, out_feats, partitioning, device, aggr="sum"): self.partitioning = partitioning self.aggr = aggr - def forward(self, gcn, graph, inputs, ampbyp=None): + 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: @@ -571,12 +573,12 @@ def forward(self, gcn, graph, inputs, ampbyp=None): 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) + return GCNFuncNONE.apply(gcn, graph, inputs, self.weight, epoch) class GCNFuncNONE(torch.autograd.Function): @staticmethod - def forward(ctx, self, graph, inputs, weight): + def forward(ctx, self, graph, inputs, weight, epoch=0): # inputs: H # graph: A # weight: W @@ -587,7 +589,7 @@ def forward(ctx, self, graph, inputs, weight): 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": @@ -599,35 +601,75 @@ def forward(ctx, self, graph, inputs, weight): ctx.save_for_backward(inputs, weight) ctx.graph = graph - ctx.self = self + # 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 + # self = ctx.self + device = ctx.device + aggr = ctx.aggr + timings = ctx.timings ag = torch.cuda.FloatTensor(graph.size(1), grad_output.size(1), \ - device=self.device).fill_(0) + 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) - - if self.aggr == "mean": + 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()) - grad_weight = outer_product(inputs.t(), ag, self.group) - - return None, None, grad_input, grad_weight + 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 diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index ddd9c51..366ffec 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -83,6 +83,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, total_start_timer.record() for i in range(n_layers): + print(f"Sampling layer {i}", flush=True) if i == 0: nnz = batch_size else: @@ -113,6 +114,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() + print(f"after gen_prob_dist", flush=True) next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, @@ -172,6 +174,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, frontiers[i + 1] = next_frontier_select.clone() current_frontier = next_frontier timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) + del next_frontier_select + del current_frontier_select # print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) if baseline_compare: @@ -179,7 +183,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - if timing: + if False and 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": v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index b834618..59b92cf 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -412,7 +412,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups start_inner_timer = torch.cuda.Event(enable_timing=True) stop_inner_timer = torch.cuda.Event(enable_timing=True) - torch.cuda.nvtx.range_push("nvtx-matc-inst") start_time(start_timer) chunk_size = math.ceil(float(mata.size(1) / (size / replication))) if True or mata.layout == torch.sparse_coo or (name == "prob" and alg == "sage"): @@ -429,14 +428,12 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) - torch.cuda.nvtx.range_pop() # matc-inst 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 - torch.cuda.nvtx.range_push("nvtx-preproc-local") start_time(start_timer) am_partid = rank_col * (size // replication ** 2) + i chunk_col_start = am_partid * chunk_size @@ -466,14 +463,12 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups "sum", \ lengths=lengths) mata_chunk_crows = torch.cuda.LongTensor(mata.size(0) + 1) - mata_chunk_crows[1:] = torch.cumsum(mata_chunk_rowcount, dim=0).long() + 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) - torch.cuda.nvtx.range_pop() # preproc-local # Determine number of nonzero columns in chunk of mata - torch.cuda.nvtx.range_push("nvtx-unique") start_time(start_timer) if mata.layout == torch.sparse_coo: nnz_cols = torch.unique(mata_chunk_indices[1,:]) @@ -481,11 +476,9 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) - torch.cuda.nvtx.range_pop() # unique # timing_dict[f"spgemm-nnzcount-{name}"].append(nnz_cols.size(0)) # Gather nnz column counts on rank q - torch.cuda.nvtx.range_push("nvtx-gather-nnzcounts") start_time(start_timer) if rank == q: nnz_cols_count_list = [] @@ -496,20 +489,16 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) - torch.cuda.nvtx.range_pop() # gather-nnzcounts # Rank q allocates recv buffer for nnz col ids - torch.cuda.nvtx.range_push("nvtx-alloc-nnzbuff") 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) - torch.cuda.nvtx.range_pop() # alloc-nnzbuff # isend/irecv nnz col ids to rank q - torch.cuda.nvtx.range_push("nvtx-send-colids") start_time(start_timer) if rank == q: recv_objs = [] @@ -522,10 +511,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups else: dist.send(nnz_cols, dst=q) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-colids-{name}", barrier=True) - torch.cuda.nvtx.range_pop() # send-nnzcols # Send selected rows count - torch.cuda.nvtx.range_push("nvtx-send-rowdata") start_time(start_timer) nnz_count = 0 if rank == q: @@ -692,14 +679,12 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups torch.cuda.synchronize() stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-send-rowdata-{name}", barrier=True) - torch.cuda.nvtx.range_pop() # send-rowcounts 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) - # torch.cuda.nvtx.range_push("nvtx-spgemm-castrecv") # start_time(start_timer) # # print(f"matb_select_recv.dtype: {matb_select_recv.dtype}") # # matb_recv_indices = matb_select_recv[:2, :].long() @@ -709,9 +694,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # # 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) - # torch.cuda.nvtx.range_pop() # spgemm-castrecv - torch.cuda.nvtx.range_push("nvtx-local-spgemm") 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 @@ -737,8 +720,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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), \ + mata_recv._values().float(), matb_recv_indices, \ + matb_recv_values.float(), 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)])) @@ -804,7 +787,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups timing_dict[f"spgemm-loc-spspmm-{name}"] = [] stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-local-spgemm-{name}", barrier=True) - torch.cuda.nvtx.range_pop() # local-spgemm # print(f"send-rowdata: {timing_dict['spgemm-send-rowdata-prob']}") # print(f"send-rownnz: {timing_dict['spgemm-send-rownnz-prob']}") @@ -827,22 +809,18 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups return matc._indices(), matc._values() # # Old sparse reduction - # torch.cuda.nvtx.range_push("nvtx-matc-coalesce") # 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) - # torch.cuda.nvtx.range_pop() # matc-coalesce # # Implement sparse allreduce w/ all_gather and padding - # torch.cuda.nvtx.range_push("nvtx-reduce-nnz") # 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_pop() # reduce-nnz # torch.cuda.nvtx.range_push("nvtx-padding") # start_time(start_timer) @@ -938,13 +916,11 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # TODO: assume n_layers=1 for now # start_time(start_timer) start_timer.record() - torch.cuda.nvtx.range_push("nvtx-probability-spgemm") 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_indices, p_num_values = dist_spgemm15D(numerator, adj_matrix, replication, rank, size, row_groups, col_groups, "prob") - torch.cuda.nvtx.range_pop() stop_timer.record() torch.cuda.synchronize() time_taken = start_timer.elapsed_time(stop_timer) @@ -952,7 +928,6 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication timing_dict["probability-spgemm"].append(time_taken) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-compute-p") p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) if name == "ladies": # p_num_values = torch.square(p_num_values).double() @@ -1002,7 +977,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # torch.Size([p.size(0), node_count_total])) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-pre-loop") frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0)).fill_(0) ones = torch.cuda.IntTensor(next_frontier._indices()[0,:].size(0)).fill_(1) @@ -1013,13 +987,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, zero_idxs = p._indices()[0, (p._values() == 0)].unique() frontier_nnz_sizes[zero_idxs] = 0 - torch.cuda.nvtx.range_pop() timing_dict["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 - torch.cuda.nvtx.range_push("nvtx-sampling") # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True @@ -1028,10 +1000,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, while underfull_minibatches: iter_count += 1 start_time(sample_start_timer) - torch.cuda.nvtx.range_push("nvtx-sampling-iter") start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-prob-rowsum") if name == "ladies": ps_p_values = torch.cumsum(p._values(), dim=0).roll(1) @@ -1047,36 +1017,28 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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 - torch.cuda.nvtx.range_pop() timing_dict["prob-rowsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-gen-darts") # dart_values = torch.cuda.DoubleTensor(n_darts * mb_count).uniform_() dart_values = torch.cuda.DoubleTensor(n_darts_col * p.size(0)).uniform_() - torch.cuda.nvtx.range_pop() timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-dart-throw") # 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)) - torch.cuda.nvtx.range_pop() timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-filter-darts") 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.size(0), p._nnz()) # n_darts * mb_count, p._nnz()) - torch.cuda.nvtx.range_pop() timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-add-to-frontier") dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) next_frontier_values = torch.logical_or( dart_hits_count, @@ -1086,11 +1048,9 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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])) - torch.cuda.nvtx.range_pop() timing_dict["add-to-frontier"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-count-samples") # sampled_count = torch.sparse.sum(next_frontier_tmp, dim=1)._values() sampled_count = sampled_count.fill_(0) sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) @@ -1099,25 +1059,20 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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) - torch.cuda.nvtx.range_pop() timing_dict["count-samples"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-preproc") overflow = torch.clamp(sampled_count - frontier_size, min=0).int() overflowed_minibatches = (overflow > 0).any() timing_dict["select-preproc"].append(stop_time(start_timer, stop_timer)) start_time(select_start_timer) - torch.cuda.nvtx.range_push("nvtx-dart-selection") if rank_col == 0 and overflowed_minibatches: while overflowed_minibatches: start_time(select_iter_start_timer) - torch.cuda.nvtx.range_push("nvtx-selection-iter") start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-psoverflow") selection_iter_count += 1 ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() @@ -1128,13 +1083,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-psoverflow"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-reciprocal") 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) - torch.cuda.nvtx.range_push("nvtx-select-instmtx") # dart_hits_inv_mtx = torch.sparse_coo_tensor(indices=next_frontier._indices(), # values=dart_hits_inv, # size=(p.size(0), node_count_total)) @@ -1143,7 +1096,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-instmtx"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-invsum") dart_hits_inv_sum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) # dart_hits_inv_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) dart_hits_inv_nnzidxs = dart_hits_inv.nonzero().squeeze() @@ -1152,13 +1104,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-invsum"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-psinv") 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) - torch.cuda.nvtx.range_push("nvtx-select-computedarts") # dart_select = torch.cuda.DoubleTensor(total_overflow).uniform_() dart_select = torch.cuda.DoubleTensor(n_darts_select).uniform_() # Compute darts for selection @@ -1168,7 +1118,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-throwdarts") # 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, @@ -1177,7 +1126,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-throwdarts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-add-to-frontier") # 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(), @@ -1188,7 +1136,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-add-to-frontier"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-samplecount") 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()) @@ -1199,7 +1146,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-samplecount"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-select-overflow") 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)) @@ -1220,7 +1166,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, dist.barrier(group=row_groups[rank_c]) - torch.cuda.nvtx.range_pop() # nvtx-dart-selection timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) # dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) @@ -1235,18 +1180,15 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, torch.Size([p.size(0), node_count_total])) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-set-probs") 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,:]) p._values()[filled_minibatches_mask] = 0 - torch.cuda.nvtx.range_pop() timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - torch.cuda.nvtx.range_push("nvtx-compute-bool") next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] sampled_count = torch.cuda.IntTensor(p.size(0)).fill_(0) @@ -1254,15 +1196,12 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ next_frontier._values()[next_frontier_nnzvals].int()) underfull_minibatches = (sampled_count < frontier_size).any() - torch.cuda.nvtx.range_pop() timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) overflow = torch.clamp(sampled_count - frontier_size, min=0).int() - torch.cuda.nvtx.range_pop() # nvtx-sampling-iter timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) - torch.cuda.nvtx.range_pop() # nvtx-sampling print(f"iter_count: {iter_count}") print(f"selection_iter_count: {selection_iter_count}") diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index f82485b..7a96096 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -13,19 +13,23 @@ import torch_geometric import torch_geometric.transforms as T from torch_geometric.data import Data, Dataset +from torch_geometric.nn import SAGEConv from torch_geometric.datasets import Planetoid, Reddit from torch_geometric.loader import NeighborSampler from torch_geometric.utils import add_remaining_self_loops import torch_sparse +import torchviz 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 import ogb from ogb.nodeproppred import PygNodePropPredDataset +from sparse_coo_tensor_cpp import sort_dst_proc_gpu import socket @@ -46,7 +50,22 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa self.replication = replication self.timings = dict() - # input layer + self.timings["total"] = [] + self.timings["sample"] = [] + self.timings["extract"] = [] + self.timings["train"] = [] + self.timings["selectfeats"] = [] + self.timings["fwd"] = [] + self.timings["bwd"] = [] + self.timings["loss"] = [] + + 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): @@ -54,15 +73,25 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa # output layer self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + # self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) + # for _ in range(n_layers - 2): + # self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) + # self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) - def forward(self, graphs, inputs): + + def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): - h = layer(self, graphs[l], h) + h = layer(self, graphs[l], h, epoch) # GCNConv + # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv + # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + # h = self.layers[l](h, edge_index) # SAGEConv if l != len(self.layers) - 1: - h = CAGF.relu(h, self.partitioning) + # h = CAGF.relu(h, self.partitioning) + h = F.relu(h) - h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + # h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + h = F.log_softmax(h, dim=1) return h @torch.no_grad() @@ -79,9 +108,37 @@ def evaluate(self, graph, features): torch.cuda.FloatTensor(edge_index.size(1)).fill_(1.0), size) adj_batch = adj_batch.t().coalesce() - h_batch = features[n_id].to(self.device) - h = layer(self, adj_batch, h_batch) + src_vtxs_sort = torch.cuda.LongTensor(n_id.size(0)) + og_idxs = torch.cuda.LongTensor(n_id.size(0)) + tally = torch.cuda.IntTensor(self.size).fill_(0) + + node_count = graph.size(1) + n_id = n_id.to(self.device) + sort_dst_proc_gpu(n_id, src_vtxs_sort, og_idxs, tally, node_count, self.size) + src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) + + output_tally = [] + for j in range(self.size): + output_tally.append(torch.cuda.IntTensor(1)) + input_tally = list(torch.split(tally, 1)) + + dist.all_to_all(output_tally, input_tally) + + 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) + output_src_vtxs -= (node_count // self.size) * self.rank + + input_features = features[output_src_vtxs] + output_features = torch.cuda.FloatTensor(n_id.size(0), features.size(1)) + dist.all_to_all_single(output_features, input_features, input_tally, output_tally) + + h_batch = torch.cuda.FloatTensor(output_features.size()) + h_batch[og_idxs] = output_features + # h_batch = features[n_id].to(self.device) + + h = layer(self, adj_batch, h_batch) # GCNConv + # h = layer(h_batch, edge_index) # SAGEConv if l != len(self.layers) - 1: h = CAGF.relu(h, self.partitioning) hs.append(h) @@ -260,7 +317,8 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep def one5d_partition_mb(rank, size, batches, replication, mb_count): rank_c = rank // replication - batch_partitions = torch.split(batches, math.ceil(float(mb_count / (size / replication))), dim=0) + # batch_partitions = torch.split(batches, math.ceil(float(mb_count / (size / replication))), dim=0) + batch_partitions = torch.split(batches, int(mb_count // (size / replication)), dim=0) return batch_partitions[rank_c] def main(args, batches=None): @@ -285,7 +343,12 @@ def main(args, batches=None): torch.cuda.set_device(rank % args.gpu) device = torch.device(f'cuda:{rank % args.gpu}') - print(f"device1: {device}", flush=True) + + 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) @@ -296,9 +359,7 @@ def main(args, batches=None): dataset = Reddit(path) data = dataset[0] - print(f"before data loading", flush=True) data = data.to(device) - print(f"after1 data loading", flush=True) data.x.requires_grad = True inputs = data.x.to(device) inputs.requires_grad = True @@ -309,19 +370,22 @@ def main(args, batches=None): adj_matrix = edge_index elif args.dataset == "Amazon": print(f"Loading coo...", flush=True) - edge_index = torch.load("../../data/Amazon/processed/data.pt") + edge_index = torch.load("/global/homes/a/alokt/data/Amazon/processed/data.pt") print(f"Done loading coo", flush=True) - n = 14249639 + # n = 14249639 + n = 9430088 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() + data.train_mask = torch.zeros(n).long() + data.train_mask[:args.batch_size * args.n_bulkmb] = 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 + # inputs.requires_grad = True data.y = data.y.to(device) elif args.dataset == "Protein": print(f"Loading coo...", flush=True) @@ -369,6 +433,7 @@ def main(args, batches=None): # 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) edge_index = data.edge_index @@ -377,6 +442,7 @@ def main(args, batches=None): adj_matrix = edge_index split_idx = dataset.get_idx_split() train_idx = split_idx['train'].to(device) + test_idx = split_idx['test'].to(device) # print(f"adj_matrix.size(): {adj_matrix.size()}") # print(f"adj_matrix: {adj_matrix}") @@ -391,9 +457,7 @@ def main(args, batches=None): partitioning = Partitioning.ONE5D - print(f"before get_proc_groups", flush=True) row_groups, col_groups = get_proc_groups(rank, size, args.replication) - print(f"after get_proc_groups", flush=True) rank_c = rank // args.replication rank_col = rank % args.replication @@ -468,18 +532,22 @@ def main(args, batches=None): print(f"recv train_idx: {train_idx}", flush=True) torch.cuda.synchronize() features_loc = inputs + features_loc.to(device) else: - _, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ - args.replication, args.normalize) + features_loc, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + args.replication, args.normalize) print("coalescing", flush=True) g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) g_loc = g_loc.to(device) g_loc = g_loc.double() - features_loc = inputs + # features_loc = inputs + features_loc = features_loc.to(device) + del inputs print("end partitioning", flush=True) print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) + print(f"features_loc.size: {features_loc.size()}", flush=True) # 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) @@ -490,22 +558,17 @@ def main(args, batches=None): n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) - # rank_train_mask = torch.split(data.train_mask, n_per_proc, dim=0)[rank_c] - rank_val_mask = torch.split(data.val_mask, n_per_proc, dim=0)[rank_c] - # rank_test_mask = torch.split(data.test_mask, n_per_proc, dim=0)[rank_c] labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] - # rank_train_nids = rank_train_mask.nonzero().squeeze() - rank_val_nids = rank_val_mask.nonzero().squeeze() - # rank_test_nids = rank_test_mask.nonzero().squeeze() if not args.dataset.startswith("ogb"): train_nid = data.train_mask.nonzero().squeeze() test_nid = data.test_mask.nonzero().squeeze() 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()}") - torch.cuda.nvtx.range_pop() - # # # adj_matrix = None # Uncomment bottom for testing # adj_matrix = adj_matrix.cuda() # adj_matrix = torch.sparse_coo_tensor(adj_matrix, @@ -522,7 +585,7 @@ def main(args, batches=None): args.n_hidden, num_classes, args.n_layers, - "mean", + args.aggr, rank, size, Partitioning.NONE, @@ -530,17 +593,17 @@ def main(args, batches=None): device, row_groups=row_groups, col_groups=col_groups) + model = model.to(device) # use optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) - for W in model.parameters(): - print(W.size()) - print(W) - dur = [] # torch.manual_seed(0) - total_start = time.time() + tally = torch.cuda.IntTensor(size).fill_(0) + rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + if rank == size - 1: + rank_n_bulkmb = args.n_bulkmb - rank_n_bulkmb * (size - 1) for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) if epoch >= 1: @@ -548,183 +611,251 @@ def main(args, batches=None): model.train() 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: - batches = torch.cuda.IntTensor(args.n_bulkmb, args.batch_size) # initially the minibatch, note row-major - torch.cuda.nvtx.range_push("nvtx-gen-minibatch-vtxs") - # torch.manual_seed(0) + torch.manual_seed(epoch) 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 = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) - - torch.cuda.nvtx.range_push("nvtx-gen-sparse-batches") - batches_indices_rows = torch.arange(batches_loc.size(0)).to(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.DoubleTensor(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) + batches_all = train_nid[vertex_perm] torch.cuda.nvtx.range_pop() - 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 * 2 - - nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) # for sa-spgemm - nnz_row_masks.fill_(0) - - nnz_recv_upperbound = edge_count // (size // args.replication) - - print("Sampling", flush=True) - if args.sample_method == "ladies": - # torch.manual_seed(rank_col) - print("first (warmup) run") - current_frontier, next_frontier, adj_matrices = \ - 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) - - print("second runs", flush=True) - nnz_row_masks.fill_(False) - # torch.cuda.profiler.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("nvtx-sampler") - # torch.manual_seed(rank_col) - 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) - elif args.sample_method == "sage": - nnz_row_masks.fill_(False) - torch.cuda.profiler.cudart().cudaProfilerStart() - torch.cuda.nvtx.range_push("nvtx-sampler") - 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, \ - args.replication, nnz_row_masks, - rank, size, row_groups, - col_groups, args.timing, args.baseline) + last_batch = False + 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 + batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) + batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) + + batches_indices_rows = torch.arange(batches_loc.size(0)).to(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.DoubleTensor(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) + + 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 * 2 + + # for sa-spgemm + nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) + nnz_row_masks.fill_(0) - torch.cuda.nvtx.range_pop() + nnz_recv_upperbound = edge_count // (size // args.replication) - print("Extracting batches", flush=True) - # adj_matrices[i][j] -- layer i mb j - rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) - adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] - frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] - for i in range(args.n_layers + 1): - for j in range(rank_n_bulkmb): - if i == 0: - row_select_min = j * args.batch_size - row_select_max = (j + 1) * args.batch_size - else: - row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) - row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) - adj_row_select_min = j * args.batch_size * (args.samp_num ** i) - adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) + if epoch >= 1: + start_time(start_timer) + print("Sampling", flush=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) + elif args.sample_method == "sage": + nnz_row_masks.fill_(False) + 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, \ + args.replication, nnz_row_masks, + rank, size, row_groups, + col_groups, args.timing, args.baseline) - if i < args.n_layers: - sampled_indices = adj_matrices_bulk[i]._indices() - sampled_values = adj_matrices_bulk[i]._values() - - sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < adj_row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - # adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_indices[0,:] -= adj_row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask].float() - - if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) + if epoch >= 1: + model.timings["sample"].append(stop_time(start_timer, stop_timer)) + torch.cuda.nvtx.range_pop() + + if epoch >= 1: + start_time(start_timer) + print("Extracting batches", flush=True) + torch.cuda.nvtx.range_push("nvtx-extracting") + # adj_matrices[i][j] -- layer i mb j + adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] + frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] + for i in range(args.n_layers + 1): + for j in range(rank_n_bulkmb): + if i == 0: + row_select_min = j * args.batch_size + row_select_max = (j + 1) * args.batch_size else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size * (args.samp_num ** i), \ - args.batch_size * (args.samp_num ** (i + 1)))) - adj_matrices[i][j] = adj_matrix_sample.coalesce() - frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] - frontiers[i][j] = frontiers_sample - - # return frontiers, adj_matrices, adj_matrix, col_groups - g_loc = g_loc.coalesce() - print("Training", flush=True) - for i in range(args.n_bulkmb): - # forward - batch_vtxs = frontiers[0][i].view(-1) - src_vtxs = frontiers[-1][i].view(-1) - features_batch = features_loc[src_vtxs] - adjs = [adj[i].coalesce() for adj in adj_matrices] - adjs.reverse() - logits = model(adjs, features_batch) - # loss = CAGF.cross_entropy(logits[train_nid], data.y[train_nid], train_nid.size(0), \ - # num_classes, partitioning, rank_c, col_groups[0], \ - # size // args.replication) - loss = F.nll_loss(logits, data.y[batch_vtxs]) - - 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]) + row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) + row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) + adj_row_select_min = j * args.batch_size * (args.samp_num ** i) + adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) + + if i < args.n_layers: + sampled_indices = adj_matrices_bulk[i]._indices() + sampled_values = adj_matrices_bulk[i]._values() + + sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < adj_row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + # adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_indices[0,:] -= adj_row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask].float() + + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size * (args.samp_num ** i), \ + args.batch_size * (args.samp_num ** (i + 1)))) + adj_matrices[i][j] = adj_matrix_sample.coalesce() + frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] + frontiers[i][j] = frontiers_sample + if epoch >= 1: + model.timings["extract"].append(stop_time(start_timer, stop_timer)) + + # if args.dataset != "Amazon" and args.dataset != "Protein": + # return frontiers, adj_matrices, adj_matrix, col_groups + if epoch >= 1: + start_time(start_timer) + g_loc = g_loc.coalesce() + torch.cuda.nvtx.range_pop() + + torch.cuda.synchronize() + print("Training", flush=True) + torch.cuda.nvtx.range_push("nvtx-training") + # for i in range(args.n_bulkmb): + for i in range(rank_n_bulkmb): + # print(f"batch {i}", flush=True) + # forward + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-selectfeats") + batch_vtxs = frontiers[0][i].view(-1) + src_vtxs = frontiers[-1][i].view(-1) + + src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) + og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) + tally.fill_(0) + + sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_count, size) + src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) + + output_tally = [] + for j in range(size): + output_tally.append(torch.cuda.IntTensor(1)) + input_tally = list(torch.split(tally, 1)) + + dist.all_to_all(output_tally, input_tally) + + 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) + output_src_vtxs -= (node_count // size) * rank + + input_features = features_loc[output_src_vtxs] + output_features = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) + dist.all_to_all_single(output_features, input_features, input_tally, output_tally) + + features_batch = torch.cuda.FloatTensor(output_features.size()) + features_batch[og_idxs] = output_features + # features_batch = output_features[og_idxs] + + # features_batch = features_loc[src_vtxs] + adjs = [adj[i] for adj in adj_matrices] + adjs.reverse() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-fwd") + logits = model(adjs, features_batch, epoch) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-loss") + loss = F.nll_loss(logits, data.y[batch_vtxs]) # GCNConv + # loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs]) # SAGEConv + optimizer.zero_grad() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-bwd") + loss.backward() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer)) + + for W in model.parameters(): + dist.all_reduce(W.grad) + torch.cuda.nvtx.range_push("nvtx-optstep") + optimizer.step() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["train"].append(stop_time(start_timer, stop_timer)) + + if last_batch: + # break + last_batch = False + args.n_bulkmb = tmp_bulkmb + args.batch_size = tmp_batch_size + torch.cuda.nvtx.range_pop() if epoch >= 1: dur.append(time.time() - epoch_start) - if rank == 0 and (epoch % 5 == 0 or epoch == args.n_epochs - 1): + if epoch >= 1 and (epoch % 5 == 0 or epoch == args.n_epochs - 1): print("Evaluating", flush=True) # acc = model.evaluate(g_loc, features_loc, labels_rank, rank_val_nids, \ # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) - out = model.evaluate(g_loc, features_loc) - res = out.argmax(dim=-1) == data.y - acc1 = int(res[data.train_mask].sum()) / int(data.train_mask.sum()) - acc2 = int(res[data.val_mask].sum()) / int(data.val_mask.sum()) - acc3 = int(res[data.test_mask].sum()) / int(data.test_mask.sum()) + acc1 = 0.0 + acc3 = 0.0 + if args.dataset != "Amazon" and ("Protein" not in args.dataset): + out = model.evaluate(g_loc, features_loc) + res = out.argmax(dim=-1) == data.y + acc1 = int(res[train_nid].sum()) / train_nid.size(0) + acc3 = int(res[test_nid].sum()) / test_nid.size(0) print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), 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 = model.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 args.timing: + sample_dur = [x / 1000 for x in model.timings["sample"]] + train_dur = [x / 1000 for x in model.timings["train"]] + selectfeats_dur = [x / 1000 for x in model.timings["selectfeats"]] + extract_dur = [x / 1000 for x in model.timings["extract"]] + fwd_dur = [x / 1000 for x in model.timings["fwd"]] + bwd_dur = [x / 1000 for x in model.timings["bwd"]] + loss_dur = [x / 1000 for x in model.timings["loss"]] + + precomp_time = sum(model.timings["precomp"]) + spmm_time = sum(model.timings["spmm"]) + gemmi_time = sum(model.timings["gemm_i"]) + gemmw_time = sum(model.timings["gemm_w"]) + aggr_time = sum(model.timings["aggr"]) + + print(f"sample: {np.sum(sample_dur)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)} feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)} loss: {np.sum(loss_dur)}") + print(f"precomp: {precomp_time} spmm: {spmm_time} gemmi: {gemmi_time} gemmw: {gemmw_time} aggr: {aggr_time}") + print(f"fwd_med: {np.median(fwd_dur)} fwd_avg: {np.average(fwd_dur)} fwd_max: {np.max(fwd_dur)}") + print(f"len(fwd_dur): {len(fwd_dur)}") + print(f"bwd_med: {np.median(bwd_dur)} bwd_avg: {np.average(bwd_dur)} bwd_max: {np.max(bwd_dur)}") + print(f"len(bwd_dur): {len(bwd_dur)}") + dist.barrier() + + total_stop = time.time() if __name__ == '__main__': parser = argparse.ArgumentParser(description='GCN') @@ -750,8 +881,8 @@ def main(args, batches=None): help="number of vertices per layer of layer-wise minibatch") 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("--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, diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index ffaa2b8..aacd7a5 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', 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.0/../../math_libs/11.0/include/", "-lcusparse"])], + extra_compile_args=["-g", "-O0", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/include/", "-lcusparse"])], cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 0fdeb60..1a0f49a 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -246,7 +246,6 @@ void spmm_gpu(const at::Tensor& A_rowindices, B.data(), // values CUDA_R_32F, // valueType CUSPARSE_ORDER_COL)); // order - // Row-major to column-major C.t_(); C.set_data(C.contiguous()); @@ -274,9 +273,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, CUSPARSE_CSRMM_ALG1, // alg &bufferSize)); // bufferSize - void* d_buffer = NULL; - // CHECK_ERROR(cudaMalloc(&d_buffer, bufferSize)); cudaMalloc(&d_buffer, bufferSize); CHECK_CUSPARSE(cusparseSpMM(handle, // handle, @@ -292,8 +289,8 @@ void spmm_gpu(const at::Tensor& A_rowindices, d_buffer)); // buffer - cudaFree(d_a_csrrows); - cudaFree(d_buffer); + // cudaFree(d_a_csrrows); + // cudaFree(d_buffer); // Column-major to row-major C.set_data(C.view({c_col, c_row})); @@ -905,7 +902,6 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); Normalize<<>>(output.data(), input.data(), index.data(), len); - CHECK_ERROR("normalize error") } __global__ void NormalizeBatch(float *output, int *input, int output_rows, int output_cols) { @@ -1154,6 +1150,65 @@ void rowselect_csr_gpu(const at::Tensor& nnz_cols, const at::Tensor& row_offsets CHECK_ERROR("rowselect csr error") } +__global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int node_count, int proc_count) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + int nodes_per_proc = node_count / proc_count; + for (int i = id; i < vtxs_count; i += stride) { + int vtx = (int) vtxs[i]; + int dst_proc = vtx / nodes_per_proc; + 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 node_count, int proc_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + int nodes_per_proc = node_count / proc_count; + for (int i = id; i < vtxs_count; i += stride) { + long vtx = vtxs[i]; + int dst_proc = vtx / nodes_per_proc; + 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 node_count, 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 (node_count == 0 || vtxs.sizes()[0] == 0) { + return; + } + + VtxTally<<>>(tally.data(), + vtxs.data(), + vtxs.sizes()[0], + node_count, + 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], + node_count, + proc_count); + CHECK_ERROR("sort_dst_proc 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"); @@ -1175,4 +1230,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From 4e96acf13429db034b0c177f1b0cce6392124753 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 6 Sep 2023 17:31:35 -0700 Subject: [PATCH 139/178] Full pipeline cora p>1,c=1 --- cagnet/nn/conv/gcn_conv.py | 12 +++++++ examples/gcn_15d.py | 67 +++++++++----------------------------- 2 files changed, 27 insertions(+), 52 deletions(-) diff --git a/cagnet/nn/conv/gcn_conv.py b/cagnet/nn/conv/gcn_conv.py index 3e1347f..6efd84b 100644 --- a/cagnet/nn/conv/gcn_conv.py +++ b/cagnet/nn/conv/gcn_conv.py @@ -583,6 +583,18 @@ def forward(ctx, self, graph, inputs, weight, epoch=0): # 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) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 7a96096..6ce7f21 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -96,48 +96,22 @@ def forward(self, graphs, inputs, epoch): @torch.no_grad() def evaluate(self, graph, features): - subgraph_loader = NeighborSampler(graph._indices(), node_idx=None, + subgraph_loader = NeighborSampler(graph, node_idx=None, sizes=[-1], batch_size=2048, 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.to(self.device) + edge_index, _, size = adj adj_batch = torch.sparse_coo_tensor(edge_index, - torch.cuda.FloatTensor(edge_index.size(1)).fill_(1.0), + torch.FloatTensor(edge_index.size(1)).fill_(1.0), size) adj_batch = adj_batch.t().coalesce() + h_batch = features[n_id] - src_vtxs_sort = torch.cuda.LongTensor(n_id.size(0)) - og_idxs = torch.cuda.LongTensor(n_id.size(0)) - tally = torch.cuda.IntTensor(self.size).fill_(0) - - node_count = graph.size(1) - n_id = n_id.to(self.device) - sort_dst_proc_gpu(n_id, src_vtxs_sort, og_idxs, tally, node_count, self.size) - src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) - - output_tally = [] - for j in range(self.size): - output_tally.append(torch.cuda.IntTensor(1)) - input_tally = list(torch.split(tally, 1)) - - dist.all_to_all(output_tally, input_tally) - - 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) - output_src_vtxs -= (node_count // self.size) * self.rank - - input_features = features[output_src_vtxs] - output_features = torch.cuda.FloatTensor(n_id.size(0), features.size(1)) - dist.all_to_all_single(output_features, input_features, input_tally, output_tally) - - h_batch = torch.cuda.FloatTensor(output_features.size()) - h_batch[og_idxs] = output_features - # h_batch = features[n_id].to(self.device) - - h = layer(self, adj_batch, h_batch) # GCNConv + h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv # h = layer(h_batch, edge_index) # SAGEConv if l != len(self.layers) - 1: h = CAGF.relu(h, self.partitioning) @@ -145,22 +119,6 @@ def evaluate(self, graph, features): features = torch.cat(hs, dim=0) return features - # logits = model(graph, features) - # model.timings = non_eval_timings # don't include evaluation timings - - # 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 get_proc_groups(rank, size, replication): rank_c = rank // replication @@ -543,12 +501,15 @@ def main(args, batches=None): g_loc = g_loc.double() # features_loc = inputs features_loc = features_loc.to(device) - del inputs + # del inputs print("end partitioning", flush=True) print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) print(f"features_loc.size: {features_loc.size()}", flush=True) + 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()) @@ -819,15 +780,17 @@ def main(args, batches=None): if epoch >= 1: dur.append(time.time() - epoch_start) - if epoch >= 1 and (epoch % 5 == 0 or epoch == args.n_epochs - 1): + if rank == 0 and epoch >= 1 and (epoch % 5 == 0 or epoch == args.n_epochs - 1): print("Evaluating", flush=True) # acc = model.evaluate(g_loc, features_loc, labels_rank, rank_val_nids, \ # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) acc1 = 0.0 acc3 = 0.0 if args.dataset != "Amazon" and ("Protein" not in args.dataset): - out = model.evaluate(g_loc, features_loc) - res = out.argmax(dim=-1) == data.y + # out = model.evaluate(g_loc, features_loc) + out = model.evaluate(adj_matrix, inputs) + datay_cpu = data.y.cpu() + res = out.argmax(dim=-1) == datay_cpu acc1 = int(res[train_nid].sum()) / train_nid.size(0) acc3 = int(res[test_nid].sum()) / test_nid.size(0) print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) From b88502ce90a93ee27d4ec47f2c5bcb7b72280896 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 7 Sep 2023 15:25:32 -0700 Subject: [PATCH 140/178] Distributed pipeline c=1 --- cagnet/samplers/sage.py | 1 - examples/gcn_15d.py | 24 ++++++------------------ sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 16 ++++++---------- 4 files changed, 13 insertions(+), 30 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 366ffec..712c804 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -114,7 +114,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() - print(f"after gen_prob_dist", flush=True) next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 6ce7f21..cd16f53 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -708,7 +708,8 @@ def main(args, batches=None): og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) tally.fill_(0) - sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_count, size) + node_per_proc = int(math.ceil(node_count / (size / args.replication))) + sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_proc) src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) output_tally = [] @@ -720,7 +721,7 @@ def main(args, batches=None): 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) - output_src_vtxs -= (node_count // size) * rank + 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)) @@ -800,22 +801,9 @@ def main(args, batches=None): train_dur = [x / 1000 for x in model.timings["train"]] selectfeats_dur = [x / 1000 for x in model.timings["selectfeats"]] extract_dur = [x / 1000 for x in model.timings["extract"]] - fwd_dur = [x / 1000 for x in model.timings["fwd"]] - bwd_dur = [x / 1000 for x in model.timings["bwd"]] - loss_dur = [x / 1000 for x in model.timings["loss"]] - - precomp_time = sum(model.timings["precomp"]) - spmm_time = sum(model.timings["spmm"]) - gemmi_time = sum(model.timings["gemm_i"]) - gemmw_time = sum(model.timings["gemm_w"]) - aggr_time = sum(model.timings["aggr"]) - - print(f"sample: {np.sum(sample_dur)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)} feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)} loss: {np.sum(loss_dur)}") - print(f"precomp: {precomp_time} spmm: {spmm_time} gemmi: {gemmi_time} gemmw: {gemmw_time} aggr: {aggr_time}") - print(f"fwd_med: {np.median(fwd_dur)} fwd_avg: {np.average(fwd_dur)} fwd_max: {np.max(fwd_dur)}") - print(f"len(fwd_dur): {len(fwd_dur)}") - print(f"bwd_med: {np.median(bwd_dur)} bwd_avg: {np.average(bwd_dur)} bwd_max: {np.max(bwd_dur)}") - print(f"len(bwd_dur): {len(bwd_dur)}") + + print(f"sample: {np.sum(sample_dur)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)}") + print(f"feats: {np.sum(selectfeats_dur)}") dist.barrier() total_stop = time.time() diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index aacd7a5..ffaa2b8 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', ext_modules=[cpp_extension.CUDAExtension('sparse_coo_tensor_cpp', ['sparse_coo_tensor.cu'], - extra_compile_args=["-g", "-O0", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/include/", "-lcusparse"])], + extra_compile_args=["-O2", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/include/", "-lcusparse"])], cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 1a0f49a..1f2e205 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1150,11 +1150,10 @@ void rowselect_csr_gpu(const at::Tensor& nnz_cols, const at::Tensor& row_offsets CHECK_ERROR("rowselect csr error") } -__global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int node_count, int proc_count) { +__global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int nodes_per_proc) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; - int nodes_per_proc = node_count / proc_count; for (int i = id; i < vtxs_count; i += stride) { int vtx = (int) vtxs[i]; int dst_proc = vtx / nodes_per_proc; @@ -1163,12 +1162,11 @@ __global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int node_c } __global__ void SortVtxs(long *vtxs, long *src_vtx_send, int *ps_proc_tally, long *og_idx, int vtxs_count, - int node_count, int proc_count) { + int nodes_per_proc) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; - int nodes_per_proc = node_count / proc_count; for (int i = id; i < vtxs_count; i += stride) { long vtx = vtxs[i]; int dst_proc = vtx / nodes_per_proc; @@ -1179,22 +1177,21 @@ __global__ void SortVtxs(long *vtxs, long *src_vtx_send, int *ps_proc_tally, lon } 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 node_count, int proc_count) { + const at::Tensor& tally, int nodes_per_proc) { int BLOCK_SIZE = 256; int BLOCK_COUNT = std::ceil(vtxs.sizes()[0] / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - if (node_count == 0 || vtxs.sizes()[0] == 0) { + if (nodes_per_proc == 0 || vtxs.sizes()[0] == 0) { return; } VtxTally<<>>(tally.data(), vtxs.data(), vtxs.sizes()[0], - node_count, - proc_count); + nodes_per_proc); auto ps_proc_tally = tally.cumsum(0, torch::kInt32).roll(1); ps_proc_tally[0] = 0; @@ -1204,8 +1201,7 @@ void sort_dst_proc_gpu(const at::Tensor& vtxs, const at::Tensor& src_vtx_sort, c ps_proc_tally.data(), og_idxs.data(), vtxs.sizes()[0], - node_count, - proc_count); + nodes_per_proc); CHECK_ERROR("sort_dst_proc error") } From a6d7befbe50d57da0723f0e808ca53749a478bf4 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 9 Sep 2023 01:51:43 -0700 Subject: [PATCH 141/178] Split minibatches across process row for c>1 --- cagnet/samplers/sage.py | 2 +- examples/gcn_15d.py | 73 +++++++++++++++++++++++++++++++++-------- 2 files changed, 61 insertions(+), 14 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 712c804..8829bbb 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -182,7 +182,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - if False and timing: + 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": v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index cd16f53..c80ed7a 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -417,9 +417,10 @@ def main(args, batches=None): row_groups, col_groups = get_proc_groups(rank, size, args.replication) - rank_c = rank // args.replication + proc_row = size // args.replication + rank_row = rank // args.replication rank_col = rank % args.replication - if rank_c >= (size // args.replication): + if rank_row >= (size // args.replication): return @@ -519,7 +520,7 @@ def main(args, batches=None): n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) - labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_c] + labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_row] if not args.dataset.startswith("ogb"): train_nid = data.train_mask.nonzero().squeeze() @@ -562,9 +563,45 @@ def main(args, batches=None): dur = [] # torch.manual_seed(0) tally = torch.cuda.IntTensor(size).fill_(0) - rank_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) + row_n_bulkmb = int(args.n_bulkmb / (size / args.replication)) if rank == size - 1: - rank_n_bulkmb = args.n_bulkmb - rank_n_bulkmb * (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) + + for epoch in range(args.n_epochs): + print(f"Epoch: {epoch}", flush=True) + if epoch >= 1: + epoch_start = time.time() + model.train() + + 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] + torch.cuda.nvtx.range_pop() + + last_batch = False + 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 + batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) + batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) + + batches_indices_rows = torch.arange(batches_loc.size(0)).to(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)) for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) if epoch >= 1: @@ -630,11 +667,11 @@ def main(args, batches=None): elif args.sample_method == "sage": nnz_row_masks.fill_(False) 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, \ - args.replication, nnz_row_masks, - rank, size, row_groups, - col_groups, args.timing, args.baseline) + args.samp_num, args.n_bulkmb, \ + args.n_layers, args.n_darts, \ + args.replication, nnz_row_masks, + rank, size, row_groups, + col_groups, args.timing, args.baseline) if epoch >= 1: model.timings["sample"].append(stop_time(start_timer, stop_timer)) @@ -657,6 +694,14 @@ def main(args, batches=None): row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) adj_row_select_min = j * args.batch_size * (args.samp_num ** i) adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) + + rank_n_bulkmb_row = int(rank_n_bulkmb / args.replication) + if rank_col == args.replication - 1: + rank_n_bulkmb_row = rank_n_bulkmb_row - rank_n_bulkmb_row * (args.replication - 1) + row_select_min += rank_col * rank_n_bulkmb_row + row_select_max += rank_col * rank_n_bulkmb_row + adj_row_select_min += rank_col * rank_n_bulkmb_row + adj_row_select_max += rank_col * rank_n_bulkmb_row if i < args.n_layers: sampled_indices = adj_matrices_bulk[i]._indices() @@ -695,6 +740,7 @@ def main(args, batches=None): print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") # for i in range(args.n_bulkmb): + print(f"rank_n_bulkmb: {rank_n_bulkmb}", flush=True) for i in range(rank_n_bulkmb): # print(f"batch {i}", flush=True) # forward @@ -703,12 +749,15 @@ def main(args, batches=None): torch.cuda.nvtx.range_push("nvtx-selectfeats") batch_vtxs = frontiers[0][i].view(-1) src_vtxs = frontiers[-1][i].view(-1) + adjs = [adj[i] for adj in adj_matrices] + adjs.reverse() src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) tally.fill_(0) - node_per_proc = int(math.ceil(node_count / (size / args.replication))) + node_per_row = int(math.ceil(node_count / (size / args.replication))) + node_per_proc = int(math.ceil(node_count / size)) sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_proc) src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) @@ -732,8 +781,6 @@ def main(args, batches=None): # features_batch = output_features[og_idxs] # features_batch = features_loc[src_vtxs] - adjs = [adj[i] for adj in adj_matrices] - adjs.reverse() torch.cuda.nvtx.range_pop() if epoch >= 1: model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) From 226314dd2415e90b51a9362462a298cf456f2af6 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 11 Sep 2023 11:42:56 -0700 Subject: [PATCH 142/178] Fixed illegal mem access error --- cagnet/samplers/sage.py | 2 +- cagnet/samplers/utils.py | 3 +++ examples/gcn_15d.py | 11 +++++------ 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 8829bbb..85b911f 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -111,7 +111,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, sa_masks, timing_dict, "sage", timing_arg) - adj_matrix = adj_matrix.to_sparse_coo() + # adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 59b92cf..1da17a0 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -519,6 +519,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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}") @@ -916,6 +918,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # TODO: assume n_layers=1 for now # start_time(start_timer) start_timer.record() + 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, diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index c80ed7a..bed9ae7 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -337,8 +337,8 @@ def main(args, batches=None): inputs = torch.rand(n, num_features) data = Data() data.y = torch.rand(n).uniform_(0, num_classes - 1).long() - data.train_mask = torch.zeros(n).long() - data.train_mask[:args.batch_size * args.n_bulkmb] = 1 + data.train_mask = torch.ones(n).long() + # data.train_mask[:args.batch_size * 3000] = 1 data.test_mask = torch.zeros(n).long() data.test_mask[args.batch_size * args.n_bulkmb:] = 1 adj_matrix = edge_index.t_() @@ -528,8 +528,8 @@ def main(args, batches=None): 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()}") + 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() @@ -602,7 +602,6 @@ def main(args, batches=None): 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)) - for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) if epoch >= 1: epoch_start = time.time() @@ -665,7 +664,7 @@ def main(args, batches=None): nnz_row_masks, rank, size, row_groups, col_groups, args.timing) elif args.sample_method == "sage": - nnz_row_masks.fill_(False) + # nnz_row_masks.fill_(False) 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, \ From 3a9d9fe5ba54ec661dab578b4a1ef7ab084bb089 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 11 Sep 2023 20:26:07 -0700 Subject: [PATCH 143/178] Fixed illegal access error w/ floors --- cagnet/samplers/utils.py | 7 ++-- examples/gcn_15d.py | 46 +++++++++++++++++++-------- sparse-extension/sparse_coo_tensor.cu | 8 ++++- 3 files changed, 43 insertions(+), 18 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 1da17a0..d04fb24 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -413,7 +413,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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 = 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: @@ -1205,8 +1206,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) - print(f"iter_count: {iter_count}") - print(f"selection_iter_count: {selection_iter_count}") + # print(f"iter_count: {iter_count}") + # print(f"selection_iter_count: {selection_iter_count}") return next_frontier diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index bed9ae7..224c13a 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -219,7 +219,8 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep normalize): 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 = math.ceil(float(node_count) / (size / replication)) + n_per_proc = node_count // (size // replication) am_partitions = None am_pbyp = None @@ -264,7 +265,13 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep 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 / replication)), dim=0) + # 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 adj_matrix_loc = am_partitions[rank_c] inputs_loc = input_partitions[rank_c] @@ -518,7 +525,8 @@ def main(args, batches=None): 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 = math.ceil(float(g_loc.size(0)) / (size / args.replication)) + n_per_proc = g_loc.size(0) // (size // args.replication) labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_row] @@ -686,22 +694,27 @@ def main(args, batches=None): for i in range(args.n_layers + 1): for j in range(rank_n_bulkmb): if i == 0: + row_select_size = args.batch_size row_select_min = j * args.batch_size row_select_max = (j + 1) * args.batch_size else: + row_select_size = args.batch_size * (args.samp_num ** (i - 1)) row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) + + adj_select_size = args.batch_size * (args.samp_num ** i) adj_row_select_min = j * args.batch_size * (args.samp_num ** i) adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) - rank_n_bulkmb_row = int(rank_n_bulkmb / args.replication) + # rank_n_bulkmb_row = int(math.ceil(rank_n_bulkmb / args.replication)) + rank_n_bulkmb_row = rank_n_bulkmb // args.replication if rank_col == args.replication - 1: - rank_n_bulkmb_row = rank_n_bulkmb_row - rank_n_bulkmb_row * (args.replication - 1) - row_select_min += rank_col * rank_n_bulkmb_row - row_select_max += rank_col * rank_n_bulkmb_row - adj_row_select_min += rank_col * rank_n_bulkmb_row - adj_row_select_max += rank_col * rank_n_bulkmb_row - + rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) + row_select_min += row_select_size * rank_col * rank_n_bulkmb_row + row_select_max += row_select_size * rank_col * rank_n_bulkmb_row + adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb_row + adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row + if i < args.n_layers: sampled_indices = adj_matrices_bulk[i]._indices() sampled_values = adj_matrices_bulk[i]._values() @@ -739,7 +752,6 @@ def main(args, batches=None): print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") # for i in range(args.n_bulkmb): - print(f"rank_n_bulkmb: {rank_n_bulkmb}", flush=True) for i in range(rank_n_bulkmb): # print(f"batch {i}", flush=True) # forward @@ -755,9 +767,13 @@ def main(args, batches=None): og_idxs = torch.cuda.LongTensor(src_vtxs.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 = 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 + print(f"node_count: {node_count} node_per_proc: {node_per_proc} node_per_row: {node_per_row}", flush=True) sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_proc) + print(f"tally.sum: {tally.sum()}", flush=True) src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) output_tally = [] @@ -769,7 +785,9 @@ def main(args, batches=None): 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) - output_src_vtxs -= node_per_proc * rank + output_src_vtxs -= node_per_row * rank_row + # output_src_vtxs -= node_per_proc * rank + print(f"output_src_vtxs.min: {output_src_vtxs.min()} max: {output_src_vtxs.max()} features_loc.size: {features_loc.size()}", flush=True) input_features = features_loc[output_src_vtxs] output_features = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 1f2e205..84435d4 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1157,6 +1157,9 @@ __global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int nodes_ for (int i = id; i < vtxs_count; i += stride) { int vtx = (int) vtxs[i]; int dst_proc = vtx / nodes_per_proc; + if (dst_proc >= 16) { + dst_proc = 15; + } atomicAdd(&proc_tally[dst_proc], 1); } } @@ -1169,7 +1172,10 @@ __global__ void SortVtxs(long *vtxs, long *src_vtx_send, int *ps_proc_tally, lon for (int i = id; i < vtxs_count; i += stride) { long vtx = vtxs[i]; - int dst_proc = vtx / nodes_per_proc; + int dst_proc = vtx / nodes_per_proc; + if (dst_proc >= 16) { + dst_proc = 15; + } int idx = atomicAdd(&ps_proc_tally[dst_proc], 1); src_vtx_send[idx] = vtx; og_idx[idx] = i; From 78a38e8d73ed0bed8b292b172d343b01acb810b9 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 11 Sep 2023 20:26:59 -0700 Subject: [PATCH 144/178] Cleanup --- examples/gcn_15d.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 224c13a..ce0c4f0 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -771,9 +771,7 @@ def main(args, batches=None): # node_per_proc = int(math.ceil(node_count / size)) node_per_row = node_count // (size // args.replication) node_per_proc = node_count // size - print(f"node_count: {node_count} node_per_proc: {node_per_proc} node_per_row: {node_per_row}", flush=True) sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_proc) - print(f"tally.sum: {tally.sum()}", flush=True) src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) output_tally = [] @@ -787,7 +785,6 @@ def main(args, batches=None): dist.all_to_all_single(output_src_vtxs, src_vtxs_sort, output_tally, input_tally) output_src_vtxs -= node_per_row * rank_row # output_src_vtxs -= node_per_proc * rank - print(f"output_src_vtxs.min: {output_src_vtxs.min()} max: {output_src_vtxs.max()} features_loc.size: {features_loc.size()}", flush=True) input_features = features_loc[output_src_vtxs] output_features = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) From fdf1be92155a0ac61355b77e3f3ef826ba11494c Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 18 Sep 2023 02:52:33 -0700 Subject: [PATCH 145/178] Uses nsparse instead of cusparse --- cagnet/samplers/ladies.py | 2 +- cagnet/samplers/sage.py | 4 +- cagnet/samplers/utils.py | 83 +- examples/gcn_15d.py | 372 +++++-- sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 1441 ++++++++++++++++++++++++- tests/test_nccl_p2p.py | 311 +++--- tests/test_nsparse.py | 28 + 8 files changed, 1953 insertions(+), 290 deletions(-) create mode 100644 tests/test_nsparse.py diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index e74dc14..076fe9a 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -9,7 +9,7 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu + scatteri_add_gpu, rowselect_coo_gpu def start_time(timer): diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 85b911f..243fdd7 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -9,7 +9,7 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ + scatteri_add_gpu, rowselect_coo_gpu, \ sparse_coo_tensor_gpu @@ -182,7 +182,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - if timing: + if False and 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": v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index d04fb24..1d1d19a 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -8,9 +8,9 @@ compute_darts_select_gpu, throw_darts_select_gpu, \ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ - scatterd_add_gpu, scatteri_add_gpu, rowselect_coo_gpu, \ + scatteri_add_gpu, rowselect_coo_gpu, \ rowselect_csr_gpu, sparse_coo_tensor_gpu, spgemm_gpu, coogeam_gpu, \ - sparse_csr_tensor_gpu + sparse_csr_tensor_gpu, nsparse_spgemm timing = True @@ -628,25 +628,42 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups timing_dict[f"spgemm-send-calls-{name}"] = [] start_time(start_inner_timer) - selected_rows_count_recv = torch.cuda.IntTensor(1) - dist.recv(selected_rows_count_recv, tag=0, src=q) + 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() + # 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) - matb_recv_lengths = matb_recv_lengths.long() + # dist.recv(matb_recv_lengths, tag=1, src=q) + 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) - matb_recv_cols = matb_recv_cols.long() + # 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()) - dist.recv(matb_recv_values, tag=3, src=q) + # 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() + + matb_recv_lengths = matb_recv_lengths.long() + matb_recv_cols = matb_recv_cols.long() # the spgemm is better with COO for SAGE prob if not (name == "prob" and alg == "sage"): @@ -710,26 +727,50 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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)) - # matb_recv = torch.sparse_csr_tensor(matb_recv_crows, matb_recv_cols, matb_recv_values.float(), - # size=(chunk_col_size, matb.size(1))) + 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() + # 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().float(), matb_recv_indices, \ - matb_recv_values.float(), 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_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) + matc_chunk = nsparse_spgemm(mata_recv.crow_indices().int(), \ + mata_recv.col_indices().int(), \ + mata_recv.values().float(), \ + matb_recv.crow_indices().int(), \ + matb_recv.col_indices().int(), \ + matb_recv.values().float(), \ + mata_recv.size(0), mata_recv.size(1), 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_recv.size(0), device=mata_recv.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: @@ -939,7 +980,9 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication elif name == "sage": # p_num_values = torch.cuda.DoubleTensor(p_num_values.size(0)).fill_(1.0) p_num_values = p_num_values.double() - scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + # scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) + # frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) + p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) # p = torch.sparse_coo_tensor(indices=p_num_indices, # values=p_num_values, # size=(numerator.size(0), node_count_total)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index ce0c4f0..334d5d8 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -18,7 +18,6 @@ from torch_geometric.loader import NeighborSampler from torch_geometric.utils import add_remaining_self_loops import torch_sparse -import torchviz from cagnet.nn.conv import GCNConv from cagnet.partitionings import Partitioning @@ -54,7 +53,14 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa self.timings["sample"] = [] self.timings["extract"] = [] self.timings["train"] = [] - self.timings["selectfeats"] = [] + # self.timings["selectfeats"] = [] + self.timings["selectfeats-precomp"] = [] + self.timings["selectfeats-sortdst"] = [] + self.timings["selectfeats-tallysplit"] = [] + self.timings["selectfeats-tallya2a"] = [] + self.timings["selectfeats-vtxsa2a"] = [] + self.timings["selectfeats-featsa2a"] = [] + self.timings["selectfeats-p2pop"] = [] self.timings["fwd"] = [] self.timings["bwd"] = [] self.timings["loss"] = [] @@ -325,9 +331,9 @@ def main(args, batches=None): data = dataset[0] data = data.to(device) - data.x.requires_grad = True + # data.x.requires_grad = True inputs = data.x.to(device) - inputs.requires_grad = True + # inputs.requires_grad = True data.y = data.y.to(device) edge_index = data.edge_index num_features = dataset.num_features @@ -339,13 +345,14 @@ def main(args, batches=None): print(f"Done loading coo", flush=True) # n = 14249639 n = 9430088 - num_features = 300 + 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[:args.batch_size * 3000] = 1 + # data.train_mask = torch.ones(n).long() + data.train_mask = torch.zeros(n).long() + data.train_mask[:args.batch_size * 4096] = 1 data.test_mask = torch.zeros(n).long() data.test_mask[args.batch_size * args.n_bulkmb:] = 1 adj_matrix = edge_index.t_() @@ -366,7 +373,7 @@ def main(args, batches=None): data.test_mask = torch.ones(n).long() adj_matrix = edge_index.t_() data = data.to(device) - inputs.requires_grad = True + # inputs.requires_grad = True data.y = data.y.to(device) elif args.dataset == "Protein_sg2": n = 17491085 @@ -378,7 +385,7 @@ def main(args, batches=None): print(f"Done loading coo", flush=True) inputs = torch.rand(n, num_features) adj_matrix = edge_index.t_() - inputs.requires_grad = True + # inputs.requires_grad = True data = Data() data.y = torch.rand(n).uniform_(0, num_classes - 1).long() data.train_mask = torch.ones(n).long() @@ -388,21 +395,24 @@ def main(args, batches=None): elif args.dataset.startswith("ogb"): if ("papers100M" in args.dataset and rank % args.gpu == 0) or "products" in args.dataset: dataset = PygNodePropPredDataset(name=args.dataset, root="../../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 = dataset[0] # data = data.to(device) # data.x.requires_grad = True # inputs = data.x.to(device) inputs = data.x data.y = data.y.squeeze().to(device) + print(f"data.y: {data.y}", flush=True) + print(f"data.y.size: {data.y.size()}", flush=True) + print(f"data.y.dtype: {data.y.dtype}", flush=True) # inputs.requires_grad = True # data.y = data.y.to(device) - edge_index = data.edge_index 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() @@ -437,46 +447,56 @@ def main(args, batches=None): # send g_loc.nnz, g_loc, train_idx.len, and train_idx for i in range(1, args.gpu): print(f"iteration i: {i}", flush=True) - _, g_loc, _ = one5d_partition(rank + i, size, inputs, adj_matrix, data, inputs, num_classes, \ - args.replication, args.normalize) + features_loc, g_loc, _ = one5d_partition(rank + i, size, inputs, adj_matrix, data, inputs, \ + num_classes, args.replication, args.normalize) print("coalescing", flush=True) g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) g_loc = g_loc.to(device) g_loc = g_loc.double() + features_loc = features_loc.to(device) edge_count = adj_matrix.size(1) dst_gpu = rank + i - g_loc_meta = torch.cuda.LongTensor(4) + g_loc_meta = torch.cuda.LongTensor(6) g_loc_meta[0] = g_loc._nnz() g_loc_meta[1] = g_loc.size(0) g_loc_meta[2] = g_loc.size(1) g_loc_meta[3] = adj_matrix.size(1) + g_loc_meta[4] = num_features + g_loc_meta[5] = num_classes dist.send(g_loc_meta, dst=dst_gpu) dist.send(g_loc._indices(), dst=dst_gpu) dist.send(g_loc._values(), dst=dst_gpu) + dist.send(features_loc, dst=dst_gpu) + dist.send(data.y, dst=dst_gpu) 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) print(f"train_idx_len: {train_idx_len} train_idx: {train_idx}", flush=True) + print(f"test_idx_len: {test_idx_len} test_idx: {test_idx}", flush=True) dist.send(train_idx, dst=dst_gpu) + dist.send(test_idx, dst=dst_gpu) - _, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + features_loc, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ args.replication, args.normalize) print("coalescing", flush=True) g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) g_loc = g_loc.to(device) + features_loc = features_loc.to(device) g_loc = g_loc.double() edge_count = adj_matrix.size(1) - del adj_matrix # Comment when testing - del inputs + # del adj_matrix # Comment when testing + # del inputs else: src_gpu = rank - (rank % args.gpu) - g_loc_meta = torch.cuda.LongTensor(4).fill_(0) + 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()) @@ -485,20 +505,39 @@ def main(args, batches=None): 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()) + print(f"g_loc_meta: {g_loc_meta} rows: {rows}", flush=True) + inputs = torch.cuda.FloatTensor(rows, num_features) + dist.recv(inputs, src=src_gpu) + + 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, \ size=torch.Size([g_loc_meta[1], g_loc_meta[2]])) 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) print(f"recv train_idx: {train_idx}", flush=True) + print(f"recv test_idx: {test_idx}", flush=True) torch.cuda.synchronize() - features_loc = inputs - features_loc.to(device) + features_loc = inputs.to(device) else: features_loc, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ args.replication, args.normalize) @@ -515,8 +554,9 @@ def main(args, batches=None): print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) print(f"features_loc.size: {features_loc.size()}", flush=True) - adj_matrix = adj_matrix.cpu() - inputs = inputs.cpu() + if rank == 0: + 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) @@ -528,16 +568,16 @@ def main(args, batches=None): # n_per_proc = math.ceil(float(g_loc.size(0)) / (size / args.replication)) n_per_proc = g_loc.size(0) // (size // args.replication) - labels_rank = torch.split(data.y, n_per_proc, dim=0)[rank_row] - 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: train_nid = train_idx test_nid = test_idx - print(f"train_nid.size: {train_nid.size()}") - print(f"test_nid.size: {test_nid.size()}") + 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() @@ -570,7 +610,8 @@ def main(args, batches=None): dur = [] # torch.manual_seed(0) - tally = torch.cuda.IntTensor(size).fill_(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) @@ -579,6 +620,7 @@ def main(args, batches=None): if rank == size - 1: rank_n_bulkmb = row_n_bulkmb - rank_n_bulkmb * (args.replication - 1) + dist.barrier() for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) if epoch >= 1: @@ -610,7 +652,6 @@ def main(args, batches=None): 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)) - print(f"Epoch: {epoch}", flush=True) if epoch >= 1: epoch_start = time.time() model.train() @@ -618,6 +659,7 @@ def main(args, batches=None): 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) + print(f"batch_count: {batch_count}", flush=True) if batches is None: torch.manual_seed(epoch) vertex_perm = torch.randperm(train_nid.size(0)) @@ -751,85 +793,175 @@ def main(args, batches=None): torch.cuda.synchronize() print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") + bulk_batch_count = args.bulk_batch_fetch + # for i in range(args.n_bulkmb): - for i in range(rank_n_bulkmb): + for b in range(0, rank_n_bulkmb, bulk_batch_count): # print(f"batch {i}", flush=True) - # forward - if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-selectfeats") - batch_vtxs = frontiers[0][i].view(-1) - src_vtxs = frontiers[-1][i].view(-1) - adjs = [adj[i] for adj in adj_matrices] - adjs.reverse() - - src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) - og_idxs = torch.cuda.LongTensor(src_vtxs.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) - src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) - - output_tally = [] - for j in range(size): - output_tally.append(torch.cuda.IntTensor(1)) - input_tally = list(torch.split(tally, 1)) - - dist.all_to_all(output_tally, input_tally) - - 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) - 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)) - dist.all_to_all_single(output_features, input_features, input_tally, output_tally) - - features_batch = torch.cuda.FloatTensor(output_features.size()) - features_batch[og_idxs] = output_features - # features_batch = output_features[og_idxs] - - # features_batch = features_loc[src_vtxs] - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) - + src_vtxs_len = frontiers[-1][b].view(-1).size(0) + og_idxs_bulk = torch.cuda.LongTensor(bulk_batch_count * src_vtxs_len) + output_features_bulk = torch.cuda.FloatTensor(bulk_batch_count * src_vtxs_len, + features_loc.size(1)) + print(f"Bulk Fetching {b} rank_n_bulkmb: {rank_n_bulkmb}", flush=True) + elem_count = 0 + ops = [] + for i in range(b, min(rank_n_bulkmb, b + bulk_batch_count)): + # forward + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-selectfeats") + batch_vtxs = frontiers[0][i].view(-1) + src_vtxs = frontiers[-1][i].view(-1) + adjs = [adj[i] for adj in adj_matrices] + adjs.reverse() + + src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) + # og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) + og_idxs = og_idxs_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] + tally.fill_(0) + if epoch >= 1: + model.timings["selectfeats-precomp"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + # 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, size) + sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_row, proc_row) + src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) + if epoch >= 1: + model.timings["selectfeats-sortdst"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + output_tally = [] + # for j in range(size): + for j in range(proc_row): + output_tally.append(torch.cuda.IntTensor(1)) + input_tally = list(torch.split(tally, 1)) + if epoch >= 1: + model.timings["selectfeats-tallysplit"].append(stop_time(start_inner_timer,stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + # dist.all_to_all(output_tally, input_tally) + dist.all_to_all(output_tally, input_tally, group=col_groups[rank_col]) + if epoch >= 1: + model.timings["selectfeats-tallya2a"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + 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) + dist.all_to_all_single(output_src_vtxs, src_vtxs_sort, output_tally, input_tally, + group=col_groups[rank_col]) + output_src_vtxs -= node_per_row * rank_row + # output_src_vtxs -= node_per_proc * rank + if epoch >= 1: + model.timings["selectfeats-vtxsa2a"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + + input_features = features_loc[output_src_vtxs] + # output_features = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) + output_features = output_features_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] + + # dist.all_to_all_single(output_features, input_features, input_tally, output_tally) + # dist.all_to_all_single(output_features, input_features, input_tally, output_tally, + # group=col_groups[rank_col]) + + input_start = 0 + output_start = 0 + for j in range(size // args.replication): + if j != rank_row: + elem_count += input_features[output_start:output_start + output_tally[j]].numel() + ops.append(dist.P2POp(dist.isend, + input_features[output_start:output_start + output_tally[j]], + j * args.replication + rank_col)) + ops.append(dist.P2POp(dist.irecv, + output_features[input_start:input_start + input_tally[j]], + j * args.replication + rank_col)) + else: + output_features[input_start:input_start + input_tally[j]] = \ + input_features[output_start:output_start + output_tally[j]] + + input_start += input_tally[j] + output_start += output_tally[j] + if epoch >= 1: + model.timings["selectfeats-p2pop"].append(stop_time(start_inner_timer, stop_inner_timer)) + + print(f"Bulk A2A", flush=True) + elem_count_tens = torch.cuda.LongTensor(1).fill_(elem_count) + dist.all_reduce(elem_count_tens) + dist.barrier() if epoch >= 1: start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-fwd") - logits = model(adjs, features_batch, epoch) - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer)) + if len(ops) > 0: + print(f"elem_count_tens: {elem_count_tens}", flush=True) + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + torch.cuda.synchronize() if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-loss") - loss = F.nll_loss(logits, data.y[batch_vtxs]) # GCNConv - # loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs]) # SAGEConv - optimizer.zero_grad() - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) - - if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-bwd") - loss.backward() - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer)) - - for W in model.parameters(): - dist.all_reduce(W.grad) - torch.cuda.nvtx.range_push("nvtx-optstep") - optimizer.step() - torch.cuda.nvtx.range_pop() + model.timings["selectfeats-featsa2a"].append(stop_time(start_inner_timer, stop_inner_timer, barrier=True)) + + dist.barrier() + print(f"FwdBwd Propagation", flush=True) + for i in range(b, b + bulk_batch_count): + # forward + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-selectfeats") + batch_vtxs = frontiers[0][i].view(-1) + src_vtxs = frontiers[-1][i].view(-1) + adjs = [adj[i] for adj in adj_matrices] + adjs.reverse() + + og_idxs = og_idxs_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] + output_features = output_features_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] + + features_batch = torch.cuda.FloatTensor(output_features.size()) + features_batch[og_idxs] = output_features + # features_batch = output_features[og_idxs] + + # features_batch = features_loc[src_vtxs] + torch.cuda.nvtx.range_pop() # nvtx-selectfeats + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-fwd") + logits = model(adjs, features_batch, epoch) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer)) + + 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]) # SAGEConv + optimizer.zero_grad() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-bwd") + loss.backward() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer)) + + for W in model.parameters(): + dist.all_reduce(W.grad) + torch.cuda.nvtx.range_push("nvtx-optstep") + optimizer.step() + torch.cuda.nvtx.range_pop() if epoch >= 1: model.timings["train"].append(stop_time(start_timer, stop_timer)) @@ -840,6 +972,7 @@ def main(args, batches=None): args.batch_size = tmp_batch_size torch.cuda.nvtx.range_pop() + 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): @@ -848,23 +981,38 @@ def main(args, batches=None): # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) acc1 = 0.0 acc3 = 0.0 + if args.timing: + sample_dur = [x / 1000 for x in model.timings["sample"]] + train_dur = [x / 1000 for x in model.timings["train"]] + # selectfeats_dur = [x / 1000 for x in model.timings["selectfeats"]] + 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"]] + + print(f"sample: {np.sum(sample_dur)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)}", flush=True) + # print(f"feats: {np.sum(selectfeats_dur)}") + print(f"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) + print(f"median: {np.median(sf_featsa2a_dur)}", flush=True) + print(f"total: {np.sum(dur)}", flush=True) if args.dataset != "Amazon" and ("Protein" not in args.dataset): # out = model.evaluate(g_loc, features_loc) out = model.evaluate(adj_matrix, inputs) datay_cpu = data.y.cpu() res = out.argmax(dim=-1) == datay_cpu + train_nid = train_nid.cpu() + test_nid = train_nid.cpu() acc1 = int(res[train_nid].sum()) / train_nid.size(0) acc3 = int(res[test_nid].sum()) / test_nid.size(0) - print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) - - if args.timing: - sample_dur = [x / 1000 for x in model.timings["sample"]] - train_dur = [x / 1000 for x in model.timings["train"]] - selectfeats_dur = [x / 1000 for x in model.timings["selectfeats"]] - extract_dur = [x / 1000 for x in model.timings["extract"]] + train_nid = train_nid.to(device) + test_nid = train_nid.to(device) + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) + else: + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), acc3), flush=True) - print(f"sample: {np.sum(sample_dur)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)}") - print(f"feats: {np.sum(selectfeats_dur)}") dist.barrier() total_stop = time.time() @@ -913,6 +1061,8 @@ def main(args, batches=None): 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, diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index ffaa2b8..a208769 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', 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.0/../../math_libs/11.0/include/", "-lcusparse"])], + extra_compile_args=["-O2", "-I/opt/nvidia/hpc_sdk/Linux_x86_64/21.11/cuda/11.0/../../math_libs/11.0/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.cu b/sparse-extension/sparse_coo_tensor.cu index 84435d4..072607d 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -7,14 +7,24 @@ #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; @@ -31,6 +41,1293 @@ using namespace at::sparse; #define CHECK_ERROR(str) \ {cudaDeviceSynchronize(); cudaError_t err; err = cudaGetLastError(); if(err!=0) {printf("ERROR %s: %d %s\n", str, err, cudaGetErrorString(err)); fflush(stdout);}} +// NSPARSE defs +#define div_round_up(a, b) ((a % b == 0)? a / b : a / b + 1) +#define WARP_BIT 5 +#define WARP 32 +#define MAX_LOCAL_THREAD_NUM 1024 +#define MAX_THREAD_BLOCK (MAX_LOCAL_THREAD_NUM / WARP) +#define BIN_NUM 7 +#define PWARP 4 +#define IMB_PWMIN 32 +#define B_PWMIN 16 +#define IMB_MIN 512 +#define B_MIN 256 +#define IMB_PW_SH_SIZE 4096 +#define B_PW_SH_SIZE 2048 +#define IMB_SH_SIZE 1024 +#define B_SH_SIZE 512 +#define HASH_SCAL 107 // Set disjoint number to COMP_SH_SIZE + +/* 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); + } + + 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); + } + + 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); + } + + __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); + } + + 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])); + torch::Tensor d_fail_count_tens = torch::zeros({1}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + torch::Tensor d_fail_perm_tens = torch::zeros({bin->bin_size[i]}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + int *d_fail_count = d_fail_count_tens.data(); + int *d_fail_perm = d_fail_perm_tens.data(); + 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)); + + torch::Tensor d_check_tens = torch::zeros({(long)table_size}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + int *d_check = d_check_tens.data(); + 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)); + + torch::Tensor d_check_tens = torch::zeros({table_size}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + torch::Tensor d_value_tens = torch::zeros({table_size}, torch::TensorOptions() + .dtype(torch::kFloat32) + .device(torch::kCUDA)); + int *d_check = d_check_tens.data(); + float *d_value = d_value_tens.data(); // assumes real type is float + + 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_searchf(double *arr, double val, int imin, int imax) { int ans = 0; @@ -403,7 +1700,6 @@ std::vector coogeam_gpu( d_c_csrrows, // csrSortedRowPtrC, C_nnz_ptr, // nnzTotalDevHostPtr, buffer)); // workspace - CHECK_ERROR("after 2nnz"); if (NULL != C_nnz_ptr) { C_nnz = *C_nnz_ptr; } else { @@ -458,6 +1754,81 @@ std::vector coogeam_gpu( 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); + + // 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); + + 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); + + /* Set bin */ + set_min_bin(&bin, n); + + // 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); + + release_bin(bin); + + 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, @@ -639,7 +2010,6 @@ std::vector spgemm_gpu( 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) { @@ -986,28 +2356,28 @@ void shift_colselect_gpu(const at::Tensor& col_shift, int nnz, int batch_size, i 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 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; @@ -1150,22 +2520,22 @@ void rowselect_csr_gpu(const at::Tensor& nnz_cols, const at::Tensor& row_offsets CHECK_ERROR("rowselect csr error") } -__global__ void VtxTally(int *proc_tally, long *vtxs, int vtxs_count, int nodes_per_proc) { +__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 >= 16) { - dst_proc = 15; + 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 nodes_per_proc, int proc_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -1173,8 +2543,8 @@ __global__ void SortVtxs(long *vtxs, long *src_vtx_send, int *ps_proc_tally, lon for (int i = id; i < vtxs_count; i += stride) { long vtx = vtxs[i]; int dst_proc = vtx / nodes_per_proc; - if (dst_proc >= 16) { - dst_proc = 15; + if (dst_proc >= proc_count) { + dst_proc = proc_count - 1; } int idx = atomicAdd(&ps_proc_tally[dst_proc], 1); src_vtx_send[idx] = vtx; @@ -1183,7 +2553,7 @@ __global__ void SortVtxs(long *vtxs, long *src_vtx_send, int *ps_proc_tally, lon } 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) { + const at::Tensor& tally, int nodes_per_proc, int proc_count) { int BLOCK_SIZE = 256; @@ -1197,7 +2567,8 @@ void sort_dst_proc_gpu(const at::Tensor& vtxs, const at::Tensor& src_vtx_sort, c VtxTally<<>>(tally.data(), vtxs.data(), vtxs.sizes()[0], - nodes_per_proc); + nodes_per_proc, + proc_count); auto ps_proc_tally = tally.cumsum(0, torch::kInt32).roll(1); ps_proc_tally[0] = 0; @@ -1207,7 +2578,8 @@ void sort_dst_proc_gpu(const at::Tensor& vtxs, const at::Tensor& src_vtx_sort, c ps_proc_tally.data(), og_idxs.data(), vtxs.sizes()[0], - nodes_per_proc); + nodes_per_proc, + proc_count); CHECK_ERROR("sort_dst_proc error") } @@ -1226,11 +2598,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("normalize_gpu", &normalize_gpu, "Normalize values in an array based on a second array"); 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("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"); } diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index e1fe394..1abe611 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -25,10 +25,7 @@ def test_nccl(args): dist.init_process_group(backend="nccl") rank = dist.get_rank() size = dist.get_world_size() - print(f"hostname: {socket.gethostname()} rank: {rank} size: {size}") - print(f"backend: {dist.get_backend()}") - print(f"rank: {rank} visible_devices: {os.environ['CUDA_VISIBLE_DEVICES']}") - torch.cuda.set_device(rank) + torch.cuda.set_device(rank % args.gpu) start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -218,106 +215,109 @@ def test_nccl(args): # 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)) + # # 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() + # if rank == 0: + # if r == run_count - 1: + # start_time(start_timer) - print(f"batch send") - 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)) + # 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) - send_ops = [] - for i in range(size - 1): - send_ops.append(dist.P2POp(dist.isend, tensors[i], i + 1)) + # 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() - 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() + # 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)) - 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 + # send_ops = [] + # for i in range(size - 1): + # send_ops.append(dist.P2POp(dist.isend, tensors[i], i + 1)) - recv_tensor = torch.cuda.FloatTensor(buffer_size) + # if rank == 0: + # if r == run_count - 1: + # start_time(start_timer) - if rank == 0 and 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() - dist.scatter(recv_tensor, tensors, src=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() - 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"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}") @@ -351,44 +351,113 @@ def test_nccl(args): # else: # recv_tensor = torch.cuda.FloatTensor(buffer_size) # dist.recv(recv_tensor, src=0) - # dist.barrier() - torch.cuda.nvtx.range_pop() - torch.cuda.profiler.cudart().cudaProfilerStop() + + # torch.cuda.nvtx.range_pop() + # torch.cuda.profiler.cudart().cudaProfilerStop() print(f"alltoall", flush=True) - buffer_size = 10 * 2**20 - tensors = [] - for i in range(size): - tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i)) - tensors[-1] = torch.cuda.FloatTensor(0).fill_(i) + 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() - 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)) + 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}") - if rank == 0: - start_time(start_timer) + print(f"p2p cross-node send", flush=True) + buffer_size = 500 * 2**20 + # buffer_size = ((66 // 4) * 2**20 // 2) // (size - 1) + 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.scatter(recv_tensor, tensors, src=0) - dist.all_to_all(recv_tensors, tensors) - torch.cuda.synchronize() + 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 rank == 0: - 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() + 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"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") + print(f"time(ms): {seconds * 1000}") + print(f"elem_count: {elem_count}") 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) From 5005af02e40f8b21fd5930520495da7edf481371 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 18 Sep 2023 19:18:33 -0700 Subject: [PATCH 146/178] Performance debugging (e.g. merging extraction) --- examples/cagnet_sageconv_torchviz | 38 +++ examples/cagnet_torchviz | 38 +++ examples/cagnet_torchviz.svg | 164 ++++++++++ examples/gcn_15d.py | 427 +++++++++++++------------- sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 14 +- 6 files changed, 455 insertions(+), 228 deletions(-) create mode 100644 examples/cagnet_sageconv_torchviz create mode 100644 examples/cagnet_torchviz create mode 100644 examples/cagnet_torchviz.svg 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/gcn_15d.py b/examples/gcn_15d.py index 334d5d8..e789d6e 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -29,6 +29,7 @@ import ogb from ogb.nodeproppred import PygNodePropPredDataset from sparse_coo_tensor_cpp import sort_dst_proc_gpu +from torchviz import make_dot import socket @@ -53,14 +54,14 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa self.timings["sample"] = [] self.timings["extract"] = [] self.timings["train"] = [] - # self.timings["selectfeats"] = [] - self.timings["selectfeats-precomp"] = [] - self.timings["selectfeats-sortdst"] = [] - self.timings["selectfeats-tallysplit"] = [] - self.timings["selectfeats-tallya2a"] = [] - self.timings["selectfeats-vtxsa2a"] = [] - self.timings["selectfeats-featsa2a"] = [] - self.timings["selectfeats-p2pop"] = [] + self.timings["selectfeats"] = [] + # self.timings["selectfeats-precomp"] = [] + # self.timings["selectfeats-sortdst"] = [] + # self.timings["selectfeats-tallysplit"] = [] + # self.timings["selectfeats-tallya2a"] = [] + # self.timings["selectfeats-vtxsa2a"] = [] + # self.timings["selectfeats-featsa2a"] = [] + # self.timings["selectfeats-p2pop"] = [] self.timings["fwd"] = [] self.timings["bwd"] = [] self.timings["loss"] = [] @@ -79,18 +80,19 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa # output layer self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + # self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) # for _ in range(n_layers - 2): # self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) # self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) - def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): h = layer(self, graphs[l], h, epoch) # GCNConv # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + # edge_index = graphs[l] # SAGEConvfake # h = self.layers[l](h, edge_index) # SAGEConv if l != len(self.layers) - 1: # h = CAGF.relu(h, self.partitioning) @@ -440,7 +442,6 @@ def main(args, batches=None): if rank_row >= (size // args.replication): return - print("start partitioning", flush=True) if args.dataset == "ogbn-papers100M" or args.dataset == "Protein_sg2": if rank % args.gpu == 0: @@ -620,6 +621,7 @@ def main(args, batches=None): if rank == size - 1: rank_n_bulkmb = row_n_bulkmb - rank_n_bulkmb * (args.replication - 1) + print(f"rank_n_bulkmb: {rank_n_bulkmb}") dist.barrier() for epoch in range(args.n_epochs): print(f"Epoch: {epoch}", flush=True) @@ -726,62 +728,62 @@ def main(args, batches=None): model.timings["sample"].append(stop_time(start_timer, stop_timer)) torch.cuda.nvtx.range_pop() - if epoch >= 1: - start_time(start_timer) - print("Extracting batches", flush=True) - torch.cuda.nvtx.range_push("nvtx-extracting") - # adj_matrices[i][j] -- layer i mb j - adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] - frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] - for i in range(args.n_layers + 1): - for j in range(rank_n_bulkmb): - if i == 0: - row_select_size = args.batch_size - row_select_min = j * args.batch_size - row_select_max = (j + 1) * args.batch_size - else: - row_select_size = args.batch_size * (args.samp_num ** (i - 1)) - row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) - row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) - - adj_select_size = args.batch_size * (args.samp_num ** i) - adj_row_select_min = j * args.batch_size * (args.samp_num ** i) - adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) - - # rank_n_bulkmb_row = int(math.ceil(rank_n_bulkmb / args.replication)) - rank_n_bulkmb_row = rank_n_bulkmb // args.replication - if rank_col == args.replication - 1: - rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) - row_select_min += row_select_size * rank_col * rank_n_bulkmb_row - row_select_max += row_select_size * rank_col * rank_n_bulkmb_row - adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb_row - adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row - - if i < args.n_layers: - sampled_indices = adj_matrices_bulk[i]._indices() - sampled_values = adj_matrices_bulk[i]._values() - - sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < adj_row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - # adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_indices[0,:] -= adj_row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask].float() - - if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) - else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size * (args.samp_num ** i), \ - args.batch_size * (args.samp_num ** (i + 1)))) - adj_matrices[i][j] = adj_matrix_sample.coalesce() - frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] - frontiers[i][j] = frontiers_sample - if epoch >= 1: - model.timings["extract"].append(stop_time(start_timer, stop_timer)) + # if epoch >= 1: + # start_time(start_timer) + # print("Extracting batches", flush=True) + # torch.cuda.nvtx.range_push("nvtx-extracting") + # # adj_matrices[i][j] -- layer i mb j + # adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] + # frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] + # for i in range(args.n_layers + 1): + # for j in range(rank_n_bulkmb): + # if i == 0: + # row_select_size = args.batch_size + # row_select_min = j * args.batch_size + # row_select_max = (j + 1) * args.batch_size + # else: + # row_select_size = args.batch_size * (args.samp_num ** (i - 1)) + # row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) + # row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) + + # adj_select_size = args.batch_size * (args.samp_num ** i) + # adj_row_select_min = j * args.batch_size * (args.samp_num ** i) + # adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) + + # # rank_n_bulkmb_row = int(math.ceil(rank_n_bulkmb / args.replication)) + # rank_n_bulkmb_row = rank_n_bulkmb // args.replication + # if rank_col == args.replication - 1: + # rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) + # row_select_min += row_select_size * rank_col * rank_n_bulkmb_row + # row_select_max += row_select_size * rank_col * rank_n_bulkmb_row + # adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb_row + # adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row + + # if i < args.n_layers: + # sampled_indices = adj_matrices_bulk[i]._indices() + # sampled_values = adj_matrices_bulk[i]._values() + + # sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ + # (sampled_indices[0,:] < adj_row_select_max) + # adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + # # adj_matrix_sample_indices[0,:] -= row_select_min + # adj_matrix_sample_indices[0,:] -= adj_row_select_min + # adj_matrix_sample_values = sampled_values[sample_select_mask].float() + + # if args.sample_method == "ladies": + # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + # adj_matrix_sample_values, \ + # (args.batch_size, args.samp_num + args.batch_size)) + # else: + # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + # adj_matrix_sample_values, \ + # (args.batch_size * (args.samp_num ** i), \ + # args.batch_size * (args.samp_num ** (i + 1)))) + # adj_matrices[i][j] = adj_matrix_sample.coalesce() + # frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] + # frontiers[i][j] = frontiers_sample + # if epoch >= 1: + # model.timings["extract"].append(stop_time(start_timer, stop_timer)) # if args.dataset != "Amazon" and args.dataset != "Protein": # return frontiers, adj_matrices, adj_matrix, col_groups @@ -795,173 +797,150 @@ def main(args, batches=None): torch.cuda.nvtx.range_push("nvtx-training") bulk_batch_count = args.bulk_batch_fetch - # for i in range(args.n_bulkmb): - for b in range(0, rank_n_bulkmb, bulk_batch_count): - # print(f"batch {i}", flush=True) - src_vtxs_len = frontiers[-1][b].view(-1).size(0) - og_idxs_bulk = torch.cuda.LongTensor(bulk_batch_count * src_vtxs_len) - output_features_bulk = torch.cuda.FloatTensor(bulk_batch_count * src_vtxs_len, - features_loc.size(1)) - print(f"Bulk Fetching {b} rank_n_bulkmb: {rank_n_bulkmb}", flush=True) - elem_count = 0 - ops = [] - for i in range(b, min(rank_n_bulkmb, b + bulk_batch_count)): - # forward - if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-selectfeats") - batch_vtxs = frontiers[0][i].view(-1) - src_vtxs = frontiers[-1][i].view(-1) - adjs = [adj[i] for adj in adj_matrices] - adjs.reverse() + # for b in range(0, rank_n_bulkmb, bulk_batch_count): + # for i in range(b, b + bulk_batch_count): + for i in range(args.n_bulkmb): + # forward + # batch_vtxs = frontiers[0][i].view(-1) + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-extracting") + + batch_select_size = args.batch_size + batch_select_min = i * args.batch_size + batch_select_max = (i + 1) * args.batch_size + + rank_n_bulkmb_row = rank_n_bulkmb // args.replication + if rank_col == args.replication - 1: + rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) + + batch_select_min += batch_select_size * rank_col * rank_n_bulkmb_row + batch_select_max += batch_select_size * rank_col * rank_n_bulkmb_row + batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) + + # src_vtxs = frontiers[-1][i].view(-1) + src_select_size = args.batch_size * (args.samp_num ** (args.n_layers - 1)) + src_select_min = i * args.batch_size * (args.samp_num ** (args.n_layers - 1)) + src_select_max = (i + 1) * args.batch_size * (args.samp_num ** (args.n_layers - 1)) + + src_select_min += src_select_size * rank_col * rank_n_bulkmb_row + src_select_max += src_select_size * rank_col * rank_n_bulkmb_row + src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) + + adjs = [None] * args.n_layers + for l in range(args.n_layers): + adj_select_size = args.batch_size * (args.samp_num ** l) + 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 += adj_select_size * rank_col * rank_n_bulkmb_row + adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row + + sampled_indices = adj_matrices_bulk[l]._indices() + sampled_values = adj_matrices_bulk[l]._values() + + sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ + (sampled_indices[0,:] < adj_row_select_max) + adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] + # adj_matrix_sample_indices[0,:] -= row_select_min + adj_matrix_sample_indices[0,:] -= adj_row_select_min + adj_matrix_sample_values = sampled_values[sample_select_mask].float() + if args.sample_method == "ladies": + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size, args.samp_num + args.batch_size)) + else: + adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ + adj_matrix_sample_values, \ + (args.batch_size * (args.samp_num ** l), \ + args.batch_size * (args.samp_num ** (l + 1)))) + adjs[args.n_layers - l - 1] = adj_matrix_sample.coalesce() + + if epoch >= 1: + model.timings["extract"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-selectfeats") + if size > 1: src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) - # og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) - og_idxs = og_idxs_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] + og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) tally.fill_(0) - if epoch >= 1: - model.timings["selectfeats-precomp"].append(stop_time(start_inner_timer, stop_inner_timer)) - if epoch >= 1: - start_time(start_inner_timer) # 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, size) - sort_dst_proc_gpu(src_vtxs, src_vtxs_sort, og_idxs, tally, node_per_row, proc_row) + # 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_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) - if epoch >= 1: - model.timings["selectfeats-sortdst"].append(stop_time(start_inner_timer, stop_inner_timer)) - if epoch >= 1: - start_time(start_inner_timer) output_tally = [] - # for j in range(size): - for j in range(proc_row): + for j in range(size): output_tally.append(torch.cuda.IntTensor(1)) input_tally = list(torch.split(tally, 1)) - if epoch >= 1: - model.timings["selectfeats-tallysplit"].append(stop_time(start_inner_timer,stop_inner_timer)) - - if epoch >= 1: - start_time(start_inner_timer) - # dist.all_to_all(output_tally, input_tally) - dist.all_to_all(output_tally, input_tally, group=col_groups[rank_col]) - if epoch >= 1: - model.timings["selectfeats-tallya2a"].append(stop_time(start_inner_timer, stop_inner_timer)) - - if epoch >= 1: - start_time(start_inner_timer) + + 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) - dist.all_to_all_single(output_src_vtxs, src_vtxs_sort, output_tally, input_tally, - group=col_groups[rank_col]) + 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 - if epoch >= 1: - model.timings["selectfeats-vtxsa2a"].append(stop_time(start_inner_timer, stop_inner_timer)) - - if epoch >= 1: - start_time(start_inner_timer) input_features = features_loc[output_src_vtxs] - # output_features = torch.cuda.FloatTensor(src_vtxs.size(0), features_loc.size(1)) - output_features = output_features_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] - - # dist.all_to_all_single(output_features, input_features, input_tally, output_tally) - # dist.all_to_all_single(output_features, input_features, input_tally, output_tally, - # group=col_groups[rank_col]) - - input_start = 0 - output_start = 0 - for j in range(size // args.replication): - if j != rank_row: - elem_count += input_features[output_start:output_start + output_tally[j]].numel() - ops.append(dist.P2POp(dist.isend, - input_features[output_start:output_start + output_tally[j]], - j * args.replication + rank_col)) - ops.append(dist.P2POp(dist.irecv, - output_features[input_start:input_start + input_tally[j]], - j * args.replication + rank_col)) - else: - output_features[input_start:input_start + input_tally[j]] = \ - input_features[output_start:output_start + output_tally[j]] - - input_start += input_tally[j] - output_start += output_tally[j] - if epoch >= 1: - model.timings["selectfeats-p2pop"].append(stop_time(start_inner_timer, stop_inner_timer)) - - print(f"Bulk A2A", flush=True) - elem_count_tens = torch.cuda.LongTensor(1).fill_(elem_count) - dist.all_reduce(elem_count_tens) - dist.barrier() + output_features = torch.cuda.FloatTensor(src_vtxs.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[og_idxs] = output_features + else: + features_batch = features_loc[src_vtxs] + + torch.cuda.nvtx.range_pop() # nvtx-selectfeats + if epoch >= 1: + model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) + + # features_batch = torch.cuda.FloatTensor(4000, 128) + # adjs[0] = torch.randint(0, 4000, (2, 4000)).cuda().long() + # adjs[1] = torch.randint(0, 600, (2, 600)).cuda().long() + if epoch >= 1: start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-fwd") + logits = model(adjs, features_batch, epoch) + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer)) - if len(ops) > 0: - print(f"elem_count_tens: {elem_count_tens}", flush=True) - reqs = dist.batch_isend_irecv(ops) - for req in reqs: - req.wait() - torch.cuda.synchronize() if epoch >= 1: - model.timings["selectfeats-featsa2a"].append(stop_time(start_inner_timer, stop_inner_timer, barrier=True)) - - dist.barrier() - print(f"FwdBwd Propagation", flush=True) - for i in range(b, b + bulk_batch_count): - # forward - if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-selectfeats") - batch_vtxs = frontiers[0][i].view(-1) - src_vtxs = frontiers[-1][i].view(-1) - adjs = [adj[i] for adj in adj_matrices] - adjs.reverse() - - og_idxs = og_idxs_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] - output_features = output_features_bulk[(i - b) * src_vtxs.size(0):(i - b + 1) * src_vtxs.size(0)] + 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]) # SAGEConv - features_batch = torch.cuda.FloatTensor(output_features.size()) - features_batch[og_idxs] = output_features - # features_batch = output_features[og_idxs] - - # features_batch = features_loc[src_vtxs] - torch.cuda.nvtx.range_pop() # nvtx-selectfeats - - if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-fwd") - logits = model(adjs, features_batch, epoch) - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer)) - - 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]) # SAGEConv - optimizer.zero_grad() - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) - - if epoch >= 1: - start_time(start_inner_timer) - torch.cuda.nvtx.range_push("nvtx-bwd") - loss.backward() - torch.cuda.nvtx.range_pop() - if epoch >= 1: - model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer)) - - for W in model.parameters(): - dist.all_reduce(W.grad) - torch.cuda.nvtx.range_push("nvtx-optstep") - optimizer.step() - torch.cuda.nvtx.range_pop() + optimizer.zero_grad() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) + + if epoch >= 1: + start_time(start_inner_timer) + torch.cuda.nvtx.range_push("nvtx-bwd") + loss.backward() + torch.cuda.nvtx.range_pop() + if epoch >= 1: + model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer)) + + for W in model.parameters(): + dist.all_reduce(W.grad) + torch.cuda.nvtx.range_push("nvtx-optstep") + optimizer.step() + torch.cuda.nvtx.range_pop() if epoch >= 1: model.timings["train"].append(stop_time(start_timer, stop_timer)) @@ -984,19 +963,27 @@ def main(args, batches=None): if args.timing: sample_dur = [x / 1000 for x in model.timings["sample"]] train_dur = [x / 1000 for x in model.timings["train"]] - # selectfeats_dur = [x / 1000 for x in model.timings["selectfeats"]] - 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"]] + 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"]] + # 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"]] print(f"sample: {np.sum(sample_dur)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)}", flush=True) - # print(f"feats: {np.sum(selectfeats_dur)}") - print(f"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) - print(f"median: {np.median(sf_featsa2a_dur)}", flush=True) + print(f"feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)}") + 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)}") + # print(f"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) + # print(f"median: {np.median(sf_featsa2a_dur)}", flush=True) print(f"total: {np.sum(dur)}", flush=True) if args.dataset != "Amazon" and ("Protein" not in args.dataset): # out = model.evaluate(g_loc, features_loc) diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index a208769..dd35dac 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', 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.0/../../math_libs/11.0/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"])], + 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.cu b/sparse-extension/sparse_coo_tensor.cu index 072607d..cb670ea 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1501,13 +1501,13 @@ void spmm_gpu(const at::Tensor& A_rowindices, clock_t start, stop; - int32_t *d_a_csrrows; - - // int devid_old = 0; - // cudaGetDevice(&devid_old); - // cudaSetDevice(devid); + // int32_t *d_a_csrrows; + // cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); + torch::Tensor A_csrrows_tens = torch::zeros({n + 1}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(torch::kCUDA)); + int32_t *d_a_csrrows = A_csrrows_tens.data(); - cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); CHECK_CUSPARSE(cusparseXcoo2csr(handle, A_rowindices.data(), nnz, @@ -1587,7 +1587,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, // cudaFree(d_a_csrrows); - // cudaFree(d_buffer); + cudaFree(d_buffer); // Column-major to row-major C.set_data(C.view({c_col, c_row})); From e7a82d838dda3a78de910da9de5702f6b8586277 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 18 Sep 2023 21:45:24 -0700 Subject: [PATCH 147/178] CSR for post-sampling extraction --- cagnet/samplers/sage.py | 3 +- examples/gcn_15d.py | 69 +++++++++++++++++++++++++---------------- 2 files changed, 45 insertions(+), 27 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 243fdd7..cde2b06 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -169,7 +169,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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_matrices[i] = adj_matrix_sample + # adj_matrices[i] = adj_matrix_sample + adj_matrices[i] = adj_matrix_sample.to_sparse_csr() frontiers[i + 1] = next_frontier_select.clone() current_frontier = next_frontier timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index e789d6e..d69753f 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -26,8 +26,6 @@ import cagnet.nn.functional as CAGF import torch.nn.functional as F -import ogb -from ogb.nodeproppred import PygNodePropPredDataset from sparse_coo_tensor_cpp import sort_dst_proc_gpu from torchviz import make_dot @@ -53,15 +51,11 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa 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["selectfeats-precomp"] = [] - # self.timings["selectfeats-sortdst"] = [] - # self.timings["selectfeats-tallysplit"] = [] - # self.timings["selectfeats-tallya2a"] = [] - # self.timings["selectfeats-vtxsa2a"] = [] - # self.timings["selectfeats-featsa2a"] = [] - # self.timings["selectfeats-p2pop"] = [] self.timings["fwd"] = [] self.timings["bwd"] = [] self.timings["loss"] = [] @@ -395,6 +389,8 @@ def main(args, batches=None): data = data.to(device) data.y = data.y.to(device) elif args.dataset.startswith("ogb"): + import ogb + 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="../../data") data = dataset[0] @@ -802,8 +798,6 @@ def main(args, batches=None): for i in range(args.n_bulkmb): # forward # batch_vtxs = frontiers[0][i].view(-1) - if epoch >= 1: - start_time(start_inner_timer) torch.cuda.nvtx.range_push("nvtx-extracting") batch_select_size = args.batch_size @@ -835,29 +829,48 @@ def main(args, batches=None): adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb_row adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row - sampled_indices = adj_matrices_bulk[l]._indices() - sampled_values = adj_matrices_bulk[l]._values() + if epoch >= 1: + start_time(start_inner_timer) + + 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] - sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ - (sampled_indices[0,:] < adj_row_select_max) - adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - # adj_matrix_sample_indices[0,:] -= row_select_min - adj_matrix_sample_indices[0,:] -= adj_row_select_min - adj_matrix_sample_values = sampled_values[sample_select_mask].float() + 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)) + + if epoch >= 1: + start_time(start_inner_timer) if args.sample_method == "ladies": adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ adj_matrix_sample_values, \ (args.batch_size, args.samp_num + args.batch_size)) else: - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ + adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, args.batch_size * (args.samp_num ** l), + 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)))) - adjs[args.n_layers - l - 1] = adj_matrix_sample.coalesce() + if epoch >= 1: + model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer)) - if epoch >= 1: - model.timings["extract"].append(stop_time(start_inner_timer, stop_inner_timer)) + if epoch >= 1: + start_time(start_inner_timer) + g adjs[args.n_layers - l - 1] = adj_matrix_sample.to_sparse_coo() + adjs[args.n_layers - l - 1] = adj_matrix_sample.coalesce() + if epoch >= 1: + model.timings["extract-coalesce"].append(stop_time(start_inner_timer, stop_inner_timer)) if epoch >= 1: start_time(start_inner_timer) @@ -977,9 +990,13 @@ def main(args, batches=None): # 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_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)} extract: {np.sum(extract_dur)} train: {np.sum(train_dur)}", flush=True) + print(f"sample: {np.sum(sample_dur)} extract-select: {np.sum(extract_select_dur)} extract-inst: {np.sum(extract_inst_dur)} extract-coalesce: {np.sum(extract_coalesce_dur)} train: {np.sum(train_dur)}", flush=True) print(f"feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)}") 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)}") # print(f"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) From ff29fea998ccc19d5ea0f832a7687b57ec73baa7 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 21 Sep 2023 23:21:12 -0700 Subject: [PATCH 148/178] Moved partitioning up, removed some Tensors in cuda extension --- cagnet/samplers/sage.py | 3 +- cagnet/samplers/utils.py | 3 + examples/gcn_15d.py | 140 ++++++++++++++------------ sparse-extension/sparse_coo_tensor.cu | 46 +++++---- 4 files changed, 106 insertions(+), 86 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index cde2b06..beb0a0b 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -76,6 +76,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, size=(mb_count * batch_size, node_count_total)) batches_expand = batches_expand.to_sparse_csr() + adj_matrix = adj_matrix.to_sparse_csr() current_frontier = batches_expand @@ -83,7 +84,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, total_start_timer.record() for i in range(n_layers): - print(f"Sampling layer {i}", flush=True) + # print(f"Sampling layer {i}", flush=True) if i == 0: nnz = batch_size else: diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 1d1d19a..73b23f0 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -481,6 +481,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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): @@ -762,6 +764,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk_cols = matc_chunk.col_indices() # matc_chunk_values = matc_chunk.values() matc_row_lens = matc_chunk_crows[1:] - matc_chunk_crows[:-1] + lt_zero_mask = matc_row_lens < 0 matc_chunk_rows = torch.repeat_interleave( torch.arange(0, mata_recv.size(0), device=mata_recv.device), matc_row_lens) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index d69753f..18ae487 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -66,28 +66,28 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa self.timings["gemm_w"] = [] self.timings["aggr"] = [] - # # input layer - self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) - # hidden layers - for i in range(n_layers - 2): - self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) - # output layer - self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) - # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) - - # self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) - # for _ in range(n_layers - 2): - # self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) - # self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) + # # # input layer + # self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # # hidden layers + # for i in range(n_layers - 2): + # self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # # output layer + # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) + # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + + self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) + for _ in range(n_layers - 2): + self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) + self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): - h = layer(self, graphs[l], h, epoch) # GCNConv - # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv - # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + # h = layer(self, graphs[l], h, epoch) # GCNConv + nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv + edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv # edge_index = graphs[l] # SAGEConvfake - # h = self.layers[l](h, edge_index) # SAGEConv + h = self.layers[l](h, edge_index) # SAGEConv if l != len(self.layers) - 1: # h = CAGF.relu(h, self.partitioning) h = F.relu(h) @@ -113,8 +113,8 @@ def evaluate(self, graph, features): adj_batch = adj_batch.t().coalesce() h_batch = features[n_id] - h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv - # h = layer(h_batch, edge_index) # SAGEConv + # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + h = layer(h_batch, edge_index) # SAGEConv if l != len(self.layers) - 1: h = CAGF.relu(h, self.partitioning) hs.append(h) @@ -227,9 +227,9 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep am_partitions = None am_pbyp = None - inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) - torch.cuda.synchronize() + # 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 @@ -280,7 +280,7 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep 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 + return inputs_loc, adj_matrix_loc, am_partitions, input_partitions def one5d_partition_mb(rank, size, batches, replication, mb_count): rank_c = rank // replication @@ -348,7 +348,7 @@ def main(args, batches=None): 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 * 4096] = 1 + 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_() @@ -365,7 +365,9 @@ def main(args, batches=None): 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.ones(n).long() + data.train_mask = torch.zeros(n).long() + data.train_mask[:args.batch_size * 2048] = 1 data.test_mask = torch.ones(n).long() adj_matrix = edge_index.t_() data = data.to(device) @@ -390,6 +392,7 @@ def main(args, batches=None): data.y = data.y.to(device) elif args.dataset.startswith("ogb"): 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="../../data") @@ -439,35 +442,44 @@ def main(args, batches=None): return print("start partitioning", flush=True) - if args.dataset == "ogbn-papers100M" or args.dataset == "Protein_sg2": + + torch.cuda.synchronize() + if args.dataset == "ogbn-papers100M" or "Protein" in args.dataset: if rank % args.gpu == 0: + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) # 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) + for i in range(1, args.gpu): print(f"iteration i: {i}", flush=True) - features_loc, g_loc, _ = one5d_partition(rank + i, size, inputs, adj_matrix, data, inputs, \ - num_classes, args.replication, args.normalize) + rank_send_row = (rank + i) // args.replication + dst_gpu = rank + i + g_send = am_partitions[rank_send_row] print("coalescing", flush=True) - g_loc = g_loc.coalesce().t_() + g_send = g_send.coalesce().t_() print("normalizing", flush=True) - g_loc = g_loc.to(device) - g_loc = g_loc.double() - features_loc = features_loc.to(device) + 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) - dst_gpu = rank + i - g_loc_meta = torch.cuda.LongTensor(6) - g_loc_meta[0] = g_loc._nnz() - g_loc_meta[1] = g_loc.size(0) - g_loc_meta[2] = g_loc.size(1) - g_loc_meta[3] = adj_matrix.size(1) - g_loc_meta[4] = num_features - g_loc_meta[5] = num_classes - dist.send(g_loc_meta, dst=dst_gpu) - dist.send(g_loc._indices(), dst=dst_gpu) - dist.send(g_loc._values(), dst=dst_gpu) - dist.send(features_loc, dst=dst_gpu) - dist.send(data.y, dst=dst_gpu) + 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(data.y.float(), dst=dst_gpu) if args.dataset == "ogbn-papers100M": train_idx_len = torch.cuda.LongTensor(1).fill_(train_idx.size(0)) @@ -479,8 +491,8 @@ def main(args, batches=None): dist.send(train_idx, dst=dst_gpu) dist.send(test_idx, dst=dst_gpu) - features_loc, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ - args.replication, args.normalize) + # features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + # args.replication, args.normalize) print("coalescing", flush=True) g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) @@ -512,7 +524,7 @@ def main(args, batches=None): inputs = torch.cuda.FloatTensor(rows, num_features) dist.recv(inputs, src=src_gpu) - data = Data() + # data = Data() data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) dist.recv(data.y, src=src_gpu) data.y = data.y.long() @@ -536,7 +548,7 @@ def main(args, batches=None): torch.cuda.synchronize() features_loc = inputs.to(device) else: - features_loc, g_loc, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ + features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ args.replication, args.normalize) print("coalescing", flush=True) g_loc = g_loc.coalesce().t_() @@ -550,6 +562,7 @@ def main(args, batches=None): print("end partitioning", flush=True) print(f"g_loc.nnz: {g_loc._nnz()}", flush=True) print(f"features_loc.size: {features_loc.size()}", flush=True) + dist.barrier() if rank == 0: adj_matrix = adj_matrix.cpu() @@ -625,7 +638,7 @@ def main(args, batches=None): epoch_start = time.time() model.train() - print("Constructing batches", flush=True) + # 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: @@ -654,10 +667,9 @@ def main(args, batches=None): epoch_start = time.time() model.train() - 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) - print(f"batch_count: {batch_count}", flush=True) + # print(f"batch_count: {batch_count}", flush=True) if batches is None: torch.manual_seed(epoch) vertex_perm = torch.randperm(train_nid.size(0)) @@ -701,7 +713,7 @@ def main(args, batches=None): if epoch >= 1: start_time(start_timer) - print("Sampling", flush=True) + # print("Sampling", flush=True) torch.cuda.nvtx.range_push("nvtx-sampling") if args.sample_method == "ladies": current_frontier, next_frontier, adj_matrices_bulk = \ @@ -789,13 +801,13 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() torch.cuda.synchronize() - print("Training", flush=True) + # print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") bulk_batch_count = args.bulk_batch_fetch # for b in range(0, rank_n_bulkmb, bulk_batch_count): # for i in range(b, b + bulk_batch_count): - for i in range(args.n_bulkmb): + for i in range(rank_n_bulkmb): # forward # batch_vtxs = frontiers[0][i].view(-1) torch.cuda.nvtx.range_push("nvtx-extracting") @@ -828,6 +840,7 @@ def main(args, batches=None): adj_row_select_max = (i + 1) * args.batch_size * (args.samp_num ** l) adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb_row adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row + adj_row_select_max = min(adj_row_select_max, adj_matrices_bulk[l].size(0)) if epoch >= 1: start_time(start_inner_timer) @@ -853,8 +866,9 @@ def main(args, batches=None): (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, args.batch_size * (args.samp_num ** l), + 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)) @@ -867,7 +881,7 @@ def main(args, batches=None): if epoch >= 1: start_time(start_inner_timer) - g adjs[args.n_layers - l - 1] = adj_matrix_sample.to_sparse_coo() + # adjs[args.n_layers - l - 1] = adj_matrix_sample.to_sparse_coo() adjs[args.n_layers - l - 1] = adj_matrix_sample.coalesce() if epoch >= 1: model.timings["extract-coalesce"].append(stop_time(start_inner_timer, stop_inner_timer)) @@ -890,7 +904,7 @@ def main(args, batches=None): src_vtx_per_proc = src_vtxs_sort.split(tally.tolist()) output_tally = [] - for j in range(size): + for j in range(proc_row): output_tally.append(torch.cuda.IntTensor(1)) input_tally = list(torch.split(tally, 1)) @@ -918,8 +932,8 @@ def main(args, batches=None): if epoch >= 1: model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) - # features_batch = torch.cuda.FloatTensor(4000, 128) - # adjs[0] = torch.randint(0, 4000, (2, 4000)).cuda().long() + # features_batch = torch.cuda.FloatTensor(3500, 128) + # adjs[0] = torch.randint(0, 3500, (2, 3500)).cuda().long() # adjs[1] = torch.randint(0, 600, (2, 600)).cuda().long() if epoch >= 1: @@ -933,8 +947,8 @@ def main(args, batches=None): 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]) # SAGEConv + # 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 optimizer.zero_grad() torch.cuda.nvtx.range_pop() @@ -1002,7 +1016,7 @@ def main(args, batches=None): # print(f"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) # print(f"median: {np.median(sf_featsa2a_dur)}", flush=True) print(f"total: {np.sum(dur)}", flush=True) - if args.dataset != "Amazon" and ("Protein" not in args.dataset): + if False and args.dataset != "Amazon" and ("Protein" not in args.dataset): # out = model.evaluate(g_loc, features_loc) out = model.evaluate(adj_matrix, inputs) datay_cpu = data.y.cpu() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index cb670ea..138c4fc 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -57,7 +57,8 @@ using namespace at::sparse; #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 107 // Set disjoint number to COMP_SH_SIZE +#define HASH_SCAL 4 // Set disjoint number to COMP_SH_SIZE /* Structure for SpGEMM */ typedef struct { @@ -89,7 +90,7 @@ void init_bin(sfBIN *bin, int M) 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); + 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)); @@ -650,6 +651,7 @@ __global__ void set_row_nz_bin_each_gl(const int *d_arpt, const int *d_acol, hash = (bcol * HASH_SCAL) % max_row_nz; adr = offset + hash; while (1) { + // printf("infinite loop key %d adr %d offset: %d hash: %d d_check[adr] %d bcol: %d max_row_nz: %d\n", key, adr, offset, hash, d_check[adr], bcol, max_row_nz); if (d_check[adr] == key) { break; } @@ -1159,17 +1161,17 @@ void set_row_nnz(int *d_arpt, int *d_acol, { 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])); - torch::Tensor d_fail_count_tens = torch::zeros({1}, torch::TensorOptions() - .dtype(torch::kInt32) - .device(torch::kCUDA)); - torch::Tensor d_fail_perm_tens = torch::zeros({bin->bin_size[i]}, torch::TensorOptions() - .dtype(torch::kInt32) - .device(torch::kCUDA)); - int *d_fail_count = d_fail_count_tens.data(); - int *d_fail_perm = d_fail_perm_tens.data(); + 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])); + // torch::Tensor d_fail_count_tens = torch::zeros({1}, torch::TensorOptions() + // .dtype(torch::kInt32) + // .device(torch::kCUDA)); + // torch::Tensor d_fail_perm_tens = torch::zeros({bin->bin_size[i]}, torch::TensorOptions() + // .dtype(torch::kInt32) + // .device(torch::kCUDA)); + // int *d_fail_count = d_fail_count_tens.data(); + // int *d_fail_perm = d_fail_perm_tens.data(); cudaMemcpy(d_fail_count, &fail_count, sizeof(int), cudaMemcpyHostToDevice); BS = 1024; GS = bin->bin_size[i]; @@ -1182,13 +1184,13 @@ void set_row_nnz(int *d_arpt, int *d_acol, 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)); + int *d_check; + checkCudaErrors(cudaMalloc((void **)&(d_check), sizeof(int) * table_size)); - torch::Tensor d_check_tens = torch::zeros({(long)table_size}, torch::TensorOptions() - .dtype(torch::kInt32) - .device(torch::kCUDA)); - int *d_check = d_check_tens.data(); + // torch::Tensor d_check_tens = torch::zeros({(long)table_size}, torch::TensorOptions() + // .dtype(torch::kInt32) + // .device(torch::kCUDA)); + // int *d_check = d_check_tens.data(); BS = 1024; GS = div_round_up(table_size, BS); init_check<<stream[i]>>>(d_check, table_size); @@ -1197,10 +1199,10 @@ void set_row_nnz(int *d_arpt, int *d_acol, (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_check); } - // cudaFree(d_fail_count); - // cudaFree(d_fail_perm); + cudaFree(d_fail_count); + cudaFree(d_fail_perm); } break; default : From 6221c873733758261712370462b12d384ff62c69 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 25 Sep 2023 19:41:54 -0700 Subject: [PATCH 149/178] Removed coo2csr and modified evaluate() for sageconv --- .gitignore | 1 + cagnet/samplers/sage.py | 6 +- cagnet/samplers/utils.py | 22 ++--- examples/gcn_15d.py | 206 +++++++++++++++++++-------------------- 4 files changed, 113 insertions(+), 122 deletions(-) diff --git a/.gitignore b/.gitignore index d2d2b55..84de964 100644 --- a/.gitignore +++ b/.gitignore @@ -33,6 +33,7 @@ tests/test_nccl tests/pytorch_geometric/ tests/torch-quiver/ tests/*.out +examples/*.batch *.qdrep *.out *.nsys-rep diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index beb0a0b..65bb4d5 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -77,14 +77,14 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand = batches_expand.to_sparse_csr() - adj_matrix = adj_matrix.to_sparse_csr() + # adj_matrix = adj_matrix.to_sparse_csr() current_frontier = batches_expand if baseline_compare: total_start_timer.record() for i in range(n_layers): - # print(f"Sampling layer {i}", flush=True) + print(f"Sampling layer {i}", flush=True) if i == 0: nnz = batch_size else: @@ -184,7 +184,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - if False and timing: + 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": v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 73b23f0..748a15d 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -727,32 +727,28 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc += matc_chunk if mata.layout == torch.sparse_csr and name == "prob" and alg == "sage": 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 = torch.sparse_csr_tensor(mata_chunk_crows, mata_chunk_cols, mata_chunk_values, + # size=(mata.size(0), chunk_col_size)) 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() - # 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().float(), matb_recv_indices, \ # matb_recv_values.float(), mata.size(0), \ # chunk_col_size, matb.size(1), coalesced=True) - matc_chunk = nsparse_spgemm(mata_recv.crow_indices().int(), \ - mata_recv.col_indices().int(), \ - mata_recv.values().float(), \ + # 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) + 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_recv.size(0), mata_recv.size(1), matb_recv.size(1)) + 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] @@ -766,7 +762,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matc_row_lens = matc_chunk_crows[1:] - matc_chunk_crows[:-1] lt_zero_mask = matc_row_lens < 0 matc_chunk_rows = torch.repeat_interleave( - torch.arange(0, mata_recv.size(0), device=mata_recv.device), + torch.arange(0, mata.size(0), device=mata.device), matc_row_lens) matc_chunk_indices = torch.stack((matc_chunk_rows, matc_chunk_cols)) @@ -1278,7 +1274,7 @@ def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ 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() + 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") diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 18ae487..db4e88c 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -98,27 +98,54 @@ def forward(self, graphs, inputs, epoch): @torch.no_grad() def evaluate(self, graph, features): + # """ GCNConv inference """ + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[-1], batch_size=2048, + # 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_batch = features[n_id] + + # # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + # h = layer(h_batch, edge_index) # SAGEConv + # if l != len(self.layers) - 1: + # h = CAGF.relu(h, self.partitioning) + # hs.append(h) + # features = torch.cat(hs, dim=0) + # return features + + """ SAGEConv inference """ subgraph_loader = NeighborSampler(graph, node_idx=None, sizes=[-1], batch_size=2048, shuffle=False) - non_eval_timings = copy.deepcopy(self.timings) - for l, layer in enumerate(self.layers): - hs = [] + + # Compute representations of nodes layer by layer, using *all* + # available edges. This leads to faster computation in contrast to + # immediately computing the final representations of each batch. + total_edges = 0 + for i in range(self.n_layers): + xs = [] for batch_size, n_id, adj in subgraph_loader: - # edge_index, _, size = adj.to(self.device) - 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_batch = features[n_id] - - # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv - h = layer(h_batch, edge_index) # SAGEConv - if l != len(self.layers) - 1: - h = CAGF.relu(h, self.partitioning) - hs.append(h) - features = torch.cat(hs, dim=0) + edge_index, _, size = adj.to(self.device) + total_edges += edge_index.size(1) + x = features[n_id].to(self.device) + # x_target = x[:size[1]] + # x = self.convs[i]((x, x_target), edge_index) + x = self.layers[i](x, edge_index) # SAGEConv + if i != self.n_layers - 1: + x = F.relu(x) + x = x[:batch_size] + xs.append(x) + features = torch.cat(xs, dim=0) + return features def get_proc_groups(rank, size, replication): @@ -395,7 +422,7 @@ def main(args, batches=None): 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="../../data") + dataset = PygNodePropPredDataset(name=args.dataset, root="/global/u1/a/alokt/data") data = dataset[0] split_idx = dataset.get_idx_split() @@ -548,6 +575,8 @@ def main(args, batches=None): 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) print("coalescing", flush=True) @@ -631,13 +660,21 @@ def main(args, batches=None): rank_n_bulkmb = row_n_bulkmb - rank_n_bulkmb * (args.replication - 1) print(f"rank_n_bulkmb: {rank_n_bulkmb}") - dist.barrier() + # g_loc = g_loc.cpu() + g_loc = g_loc.to_sparse_csr() + g_loc = g_loc.to(device) for epoch in range(args.n_epochs): 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 = [] + # 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) @@ -645,8 +682,9 @@ def main(args, batches=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_pop() # construct-batches + torch.cuda.nvtx.range_push("nvtx-extract-batches") last_batch = False for b in range(0, batch_count, args.n_bulkmb): if b + args.n_bulkmb > batch_count: @@ -666,16 +704,18 @@ def main(args, batches=None): if epoch >= 1: epoch_start = time.time() model.train() - - 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_pop() # extract-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 for b in range(0, batch_count, args.n_bulkmb): if b + args.n_bulkmb > batch_count: @@ -685,6 +725,7 @@ def main(args, batches=None): tmp_batch_size = args.batch_size args.n_bulkmb = 1 args.batch_size = batch_count - b + torch.cuda.nvtx.range_push("nvtx-partition-batches") batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) @@ -694,6 +735,7 @@ def main(args, batches=None): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) batches_values = torch.cuda.DoubleTensor(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))) + torch.cuda.nvtx.range_pop() # partition-batches # g_loc = torch.pow(g_loc, 2) node_count = g_loc.size(1) @@ -706,7 +748,8 @@ def main(args, batches=None): args.n_darts = avg_degree * 2 # for sa-spgemm - nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) + # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) + 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) @@ -734,73 +777,12 @@ def main(args, batches=None): if epoch >= 1: model.timings["sample"].append(stop_time(start_timer, stop_timer)) - torch.cuda.nvtx.range_pop() - - # if epoch >= 1: - # start_time(start_timer) - # print("Extracting batches", flush=True) - # torch.cuda.nvtx.range_push("nvtx-extracting") - # # adj_matrices[i][j] -- layer i mb j - # adj_matrices = [[None] * rank_n_bulkmb for x in range(args.n_layers)] - # frontiers = [[None] * rank_n_bulkmb for x in range(args.n_layers + 1)] - # for i in range(args.n_layers + 1): - # for j in range(rank_n_bulkmb): - # if i == 0: - # row_select_size = args.batch_size - # row_select_min = j * args.batch_size - # row_select_max = (j + 1) * args.batch_size - # else: - # row_select_size = args.batch_size * (args.samp_num ** (i - 1)) - # row_select_min = j * args.batch_size * (args.samp_num ** (i - 1)) - # row_select_max = (j + 1) * args.batch_size * (args.samp_num ** (i - 1)) - - # adj_select_size = args.batch_size * (args.samp_num ** i) - # adj_row_select_min = j * args.batch_size * (args.samp_num ** i) - # adj_row_select_max = (j + 1) * args.batch_size * (args.samp_num ** i) - - # # rank_n_bulkmb_row = int(math.ceil(rank_n_bulkmb / args.replication)) - # rank_n_bulkmb_row = rank_n_bulkmb // args.replication - # if rank_col == args.replication - 1: - # rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) - # row_select_min += row_select_size * rank_col * rank_n_bulkmb_row - # row_select_max += row_select_size * rank_col * rank_n_bulkmb_row - # adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb_row - # adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row - - # if i < args.n_layers: - # sampled_indices = adj_matrices_bulk[i]._indices() - # sampled_values = adj_matrices_bulk[i]._values() - - # sample_select_mask = (adj_row_select_min <= sampled_indices[0,:]) & \ - # (sampled_indices[0,:] < adj_row_select_max) - # adj_matrix_sample_indices = sampled_indices[:, sample_select_mask] - # # adj_matrix_sample_indices[0,:] -= row_select_min - # adj_matrix_sample_indices[0,:] -= adj_row_select_min - # adj_matrix_sample_values = sampled_values[sample_select_mask].float() - - # if args.sample_method == "ladies": - # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - # adj_matrix_sample_values, \ - # (args.batch_size, args.samp_num + args.batch_size)) - # else: - # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - # adj_matrix_sample_values, \ - # (args.batch_size * (args.samp_num ** i), \ - # args.batch_size * (args.samp_num ** (i + 1)))) - # adj_matrices[i][j] = adj_matrix_sample.coalesce() - # frontiers_sample = frontiers_bulk[i][row_select_min:row_select_max,:] - # frontiers[i][j] = frontiers_sample - # if epoch >= 1: - # model.timings["extract"].append(stop_time(start_timer, stop_timer)) - - # if args.dataset != "Amazon" and args.dataset != "Protein": - # return frontiers, adj_matrices, adj_matrix, col_groups + torch.cuda.nvtx.range_pop() # nvtx-sampling + if epoch >= 1: start_time(start_timer) - g_loc = g_loc.coalesce() - torch.cuda.nvtx.range_pop() + # g_loc = g_loc.coalesce() - torch.cuda.synchronize() # print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") bulk_batch_count = args.bulk_batch_fetch @@ -886,6 +868,7 @@ def main(args, batches=None): if epoch >= 1: model.timings["extract-coalesce"].append(stop_time(start_inner_timer, stop_inner_timer)) + torch.cuda.nvtx.range_pop() # nvtx-extracting if epoch >= 1: start_time(start_inner_timer) torch.cuda.nvtx.range_push("nvtx-selectfeats") @@ -932,15 +915,20 @@ def main(args, batches=None): if epoch >= 1: model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) - # features_batch = torch.cuda.FloatTensor(3500, 128) - # adjs[0] = torch.randint(0, 3500, (2, 3500)).cuda().long() - # adjs[1] = torch.randint(0, 600, (2, 600)).cuda().long() + # features_batch = torch.cuda.FloatTensor(69500, 128) # SAGEConvfake + # adjs[0] = torch.randint(0, 3500, (2, 3500)).cuda().long() # SAGEConvfake + # adjs[1] = torch.randint(0, 600, (2, 600)).cuda().long() # SAGEConvfake if epoch >= 1: start_time(start_inner_timer) 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()) logits = model(adjs, features_batch, epoch) - torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() # nvtx-fwd if epoch >= 1: model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer)) @@ -951,7 +939,7 @@ def main(args, batches=None): loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs].long()) # SAGEConv optimizer.zero_grad() - torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() # nvtx-loss if epoch >= 1: model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) @@ -959,7 +947,7 @@ def main(args, batches=None): start_time(start_inner_timer) torch.cuda.nvtx.range_push("nvtx-bwd") loss.backward() - torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() # nvtx-bwd if epoch >= 1: model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer)) @@ -967,7 +955,7 @@ def main(args, batches=None): dist.all_reduce(W.grad) torch.cuda.nvtx.range_push("nvtx-optstep") optimizer.step() - torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() # nvtx-optstep if epoch >= 1: model.timings["train"].append(stop_time(start_timer, stop_timer)) @@ -976,7 +964,8 @@ def main(args, batches=None): last_batch = False args.n_bulkmb = tmp_bulkmb args.batch_size = tmp_batch_size - torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_pop() # nvtx-training + torch.cuda.nvtx.range_pop() # nvtx-bulk dist.barrier() if epoch >= 1: @@ -1013,14 +1002,19 @@ def main(args, batches=None): print(f"sample: {np.sum(sample_dur)} extract-select: {np.sum(extract_select_dur)} extract-inst: {np.sum(extract_inst_dur)} extract-coalesce: {np.sum(extract_coalesce_dur)} train: {np.sum(train_dur)}", flush=True) print(f"feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)}") 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)}") + 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"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) # print(f"median: {np.median(sf_featsa2a_dur)}", flush=True) print(f"total: {np.sum(dur)}", flush=True) - if False and args.dataset != "Amazon" and ("Protein" not in args.dataset): + if args.dataset != "Amazon" and ("Protein" not in args.dataset): # out = model.evaluate(g_loc, features_loc) + print(f"adj_matrix.device: {adj_matrix.device} inputs.device: {inputs.device}", flush=True) out = model.evaluate(adj_matrix, inputs) - datay_cpu = data.y.cpu() - res = out.argmax(dim=-1) == datay_cpu + res = out.argmax(dim=-1) == data.y train_nid = train_nid.cpu() test_nid = train_nid.cpu() acc1 = int(res[train_nid].sum()) / train_nid.size(0) From d5f582f00835014e7a85fea2632fbf20ed57e333 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 28 Sep 2023 02:50:44 -0700 Subject: [PATCH 150/178] Removed nsparse and csrinst calls --- cagnet/samplers/utils.py | 67 +++++-- examples/gcn_15d.py | 263 +++++++++++++++++--------- sparse-extension/sparse_coo_tensor.cu | 45 +++++ 3 files changed, 268 insertions(+), 107 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 748a15d..d40e597 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -10,7 +10,7 @@ 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 + sparse_csr_tensor_gpu, nsparse_spgemm, rearrange_rows_gpu timing = True @@ -584,8 +584,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, 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_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) + 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[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, \ @@ -676,8 +682,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, 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_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) + 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, dim=0) + matb_recv_crows[0] = 0 nnz_count += 2 * matb_recv_cols.size(0) + matb_recv_lengths.size(0) @@ -729,29 +741,48 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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)) - 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() + 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) - 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] + + # 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[mata_recv_rows] = matb_recv_row_lens[mata_chunk_cols] + 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) + rearrange_rows_gpu(mata_recv_rows, mata_chunk_cols, matc_chunk_crows, + matb_recv_crows, matb_recv_cols, 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() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index db4e88c..8bd6c3c 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -27,7 +27,7 @@ import torch.nn.functional as F from sparse_coo_tensor_cpp import sort_dst_proc_gpu -from torchviz import make_dot +import quiver import socket @@ -59,6 +59,7 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa self.timings["fwd"] = [] self.timings["bwd"] = [] self.timings["loss"] = [] + self.timings["fakemats"] = [] self.timings["precomp"] = [] self.timings["spmm"] = [] @@ -66,28 +67,29 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa 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)) + # input layer + self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # hidden layers + for i in range(n_layers - 2): + self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # output layer + self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) - self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) - for _ in range(n_layers - 2): - self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) - self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) + # self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) + # for _ in range(n_layers - 2): + # self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) + # self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): - # h = layer(self, graphs[l], h, epoch) # GCNConv - nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv - edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv + # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv # edge_index = graphs[l] # SAGEConvfake - h = self.layers[l](h, edge_index) # SAGEConv + # edge_index, _, size = graphs[l] # quiver + # h = self.layers[l](h, edge_index) # SAGEConv + h = layer(self, graphs[l], h, epoch) # GCNConv if l != len(self.layers) - 1: # h = CAGF.relu(h, self.partitioning) h = F.relu(h) @@ -97,57 +99,100 @@ def forward(self, graphs, inputs, epoch): return h @torch.no_grad() - def evaluate(self, graph, features): - # """ GCNConv inference """ + def evaluate(self, graph, features, test_idx, labels): # subgraph_loader = NeighborSampler(graph, node_idx=None, # sizes=[-1], batch_size=2048, - # shuffle=False) - # non_eval_timings = copy.deepcopy(self.timings) - # for l, layer in enumerate(self.layers): - # hs = [] + # 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) - # 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_batch = features[n_id] + # 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) - # # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv - # h = layer(h_batch, edge_index) # SAGEConv - # if l != len(self.layers) - 1: - # h = CAGF.relu(h, self.partitioning) - # hs.append(h) - # features = torch.cat(hs, dim=0) # return features - """ SAGEConv inference """ subgraph_loader = NeighborSampler(graph, node_idx=None, sizes=[-1], batch_size=2048, + # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, + # sizes=[-1, -1], batch_size=512, shuffle=False) - - # Compute representations of nodes layer by layer, using *all* - # available edges. This leads to faster computation in contrast to - # immediately computing the final representations of each batch. - total_edges = 0 - for i in range(self.n_layers): - xs = [] + 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) - total_edges += edge_index.size(1) - x = features[n_id].to(self.device) - # x_target = x[:size[1]] - # x = self.convs[i]((x, x_target), edge_index) - x = self.layers[i](x, edge_index) # SAGEConv - if i != self.n_layers - 1: - x = F.relu(x) - x = x[:batch_size] - xs.append(x) - features = torch.cat(xs, dim=0) - + # 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_batch = features[n_id] + + h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + hs.append(h) # GCNConv + # hs.append(h[:batch_size]) # SAGEConv + features = torch.cat(hs, dim=0) return features + # print(f"test_idx: {test_idx}", flush=True) + # non_eval_timings = copy.deepcopy(self.timings) + # correct_count = 0 + # for batch_size, n_id, adjs in subgraph_loader: + # h = features[n_id] + # for l, (edge_index, _, size) in enumerate(adjs): + # # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + # h = self.layers[l](h, edge_index) # SAGEConv + # if l != len(self.layers) - 1: + # h = F.relu(h) + # h = F.log_softmax(h, dim=1) + # batch = n_id[:batch_size] + # print(f"batch: {batch}", flush=True) + # print(f"h[:batch_size]: {h[:batch_size]}", flush=True) + # print(f"h[:batch_size].argmax(dim=-1): {h[:batch_size].argmax(dim=-1)}", flush=True) + # preds = h[:batch_size].argmax(dim=-1) == labels[batch] + # correct_count += preds.sum() + # return correct_count + + # """ SAGEConv inference """ + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[-1], batch_size=2048, + # shuffle=False) + + # # Compute representations of nodes layer by layer, using *all* + # # available edges. This leads to faster computation in contrast to + # # immediately computing the final representations of each batch. + # total_edges = 0 + # for i in range(self.n_layers): + # xs = [] + # for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + # total_edges += edge_index.size(1) + # x = features[n_id].to(self.device) + # # x_target = x[:size[1]] + # # x = self.convs[i]((x, x_target), edge_index) + # x = self.layers[i](x, edge_index) # SAGEConv + # if i != self.n_layers - 1: + # x = F.relu(x) + # x = x[:batch_size] + # xs.append(x) + # features = torch.cat(xs, dim=0) + + # return features + def get_proc_groups(rank, size, replication): rank_c = rank // replication @@ -663,7 +708,14 @@ def main(args, batches=None): # g_loc = g_loc.cpu() g_loc = g_loc.to_sparse_csr() 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() @@ -674,14 +726,15 @@ def main(args, batches=None): 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] + # 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() # construct-batches torch.cuda.nvtx.range_push("nvtx-extract-batches") @@ -703,7 +756,6 @@ def main(args, batches=None): batches_indices = torch.stack((batches_indices_rows, batches_indices_cols)) if epoch >= 1: epoch_start = time.time() - model.train() torch.cuda.nvtx.range_pop() # extract-batches # torch.cuda.nvtx.range_push("nvtx-construct-batches") @@ -785,10 +837,7 @@ def main(args, batches=None): # print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") - bulk_batch_count = args.bulk_batch_fetch - # for b in range(0, rank_n_bulkmb, bulk_batch_count): - # for i in range(b, b + bulk_batch_count): for i in range(rank_n_bulkmb): # forward # batch_vtxs = frontiers[0][i].view(-1) @@ -814,6 +863,7 @@ def main(args, batches=None): src_select_min += src_select_size * rank_col * rank_n_bulkmb_row src_select_max += src_select_size * rank_col * rank_n_bulkmb_row src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) + # print(f"src_vtxs: {src_vtxs} src_vtxs.size: {src_vtxs.size()} zero_count: {(src_vtxs == 0).sum()}", flush=True) adjs = [None] * args.n_layers for l in range(args.n_layers): @@ -858,6 +908,8 @@ def main(args, batches=None): adj_sample_values, \ (args.batch_size * (args.samp_num ** l), \ args.batch_size * (args.samp_num ** (l + 1)))) + adj_nnzs.append(adj_sample_values.size(0)) + # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") if epoch >= 1: model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer)) @@ -915,18 +967,48 @@ def main(args, batches=None): if epoch >= 1: model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) - # features_batch = torch.cuda.FloatTensor(69500, 128) # SAGEConvfake - # adjs[0] = torch.randint(0, 3500, (2, 3500)).cuda().long() # SAGEConvfake - # adjs[1] = torch.randint(0, 600, (2, 600)).cuda().long() # SAGEConvfake + # if epoch >= 1: + # start_time(start_inner_timer) + # torch.cuda.nvtx.range_push("nvtx-fakemats") + # features_batch = torch.cuda.FloatTensor(69740, 128) # SAGEConvfake + # adjs0_rows = torch.randint(0, 9862, (91518,)).cuda().long() # SAGEConvfake + # adjs0_cols = torch.randint(0, 69740, (91518,)).cuda().long() # SAGEConvfake + # adjs[0] = torch.stack((adjs0_rows, adjs0_cols)) # SAGEConvfake + # adjs1_rows = torch.randint(0, 1024, (9065,)).cuda().long() # SAGEConvfake + # adjs1_cols = torch.randint(0, 9862, (9065,)).cuda().long() # SAGEConvfake + # adjs[1] = torch.stack((adjs1_rows, adjs1_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) 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()) + # 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: @@ -935,10 +1017,11 @@ def main(args, batches=None): 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 + print(f"logits: {logits[:args.batch_size]}", flush=True) + loss = F.nll_loss(logits, data.y[batch_vtxs].long()) # GCNConv + # loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs].long()) # SAGEConv + # loss = F.nll_loss(logits[:args.batch_size], data.y[seeds].long()) # quiver - optimizer.zero_grad() torch.cuda.nvtx.range_pop() # nvtx-loss if epoch >= 1: model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) @@ -987,6 +1070,7 @@ def main(args, batches=None): 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"]] @@ -1001,26 +1085,27 @@ def main(args, batches=None): print(f"sample: {np.sum(sample_dur)} extract-select: {np.sum(extract_select_dur)} extract-inst: {np.sum(extract_inst_dur)} extract-coalesce: {np.sum(extract_coalesce_dur)} train: {np.sum(train_dur)}", flush=True) print(f"feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)}") - 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)}") - 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"precomp: {np.sum(sf_precomp_dur)} sortdst: {np.sum(sf_sortdst_dur)} tallysplit: {np.sum(sf_tallysplit_dur)} tallya2a: {np.sum(sf_tallya2a_dur)} vtxsa2a: {np.sum(sf_vtxsa2a_dur)} featsa2a: {np.sum(sf_featsa2a_dur)}", flush=True) - # print(f"median: {np.median(sf_featsa2a_dur)}", flush=True) + 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)}", flush=True) if args.dataset != "Amazon" and ("Protein" not in args.dataset): - # out = model.evaluate(g_loc, features_loc) - print(f"adj_matrix.device: {adj_matrix.device} inputs.device: {inputs.device}", flush=True) - out = model.evaluate(adj_matrix, inputs) - res = out.argmax(dim=-1) == data.y + model = model.cpu() train_nid = train_nid.cpu() - test_nid = train_nid.cpu() - acc1 = int(res[train_nid].sum()) / train_nid.size(0) + test_nid = test_nid.cpu() + out = model.evaluate(adj_matrix, inputs, test_nid, data.y) + # correct_count = model.evaluate(adj_matrix, inputs, test_nid, data.y.cpu()) + res = out.argmax(dim=-1) == data.y.cpu() + # acc1 = int(res[train_nid].sum()) / train_nid.size(0) acc3 = int(res[test_nid].sum()) / test_nid.size(0) + # acc3 = correct_count / test_nid.size(0) train_nid = train_nid.to(device) - test_nid = train_nid.to(device) + test_nid = test_nid.to(device) + model = model.to(device) print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) else: print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), acc3), flush=True) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 138c4fc..8badd94 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -8,6 +8,7 @@ #include #include +#include #include #include "cusparse.h" @@ -2585,6 +2586,49 @@ void sort_dst_proc_gpu(const at::Tensor& vtxs, const at::Tensor& src_vtx_sort, c CHECK_ERROR("sort_dst_proc error") } +__global__ void RearrangeRows(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 rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols, const at::Tensor& matc_crows, + const at::Tensor& matb_crows, const at::Tensor& matb_cols, + const at::Tensor& matc_cols) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(mata_rows.sizes()[0] / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (mata_rows.sizes()[0] == 0) { + return; + } + + RearrangeRows<<>>(mata_rows.data(), + mata_cols.data(), + matc_crows.data(), + matb_crows.data(), + matb_cols.data(), + matc_cols.data(), + mata_rows.sizes()[0]); + + CHECK_ERROR("rearrange_rows error") +} + 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"); @@ -2608,4 +2652,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From acaeb6df144be0452a8fc8eb13dddb65c31f3678 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 2 Oct 2023 20:30:05 -0700 Subject: [PATCH 151/178] Kind of added dataloader and fixed ranks training the same batch --- cagnet/samplers/sage.py | 50 +++--- cagnet/samplers/utils.py | 25 ++- examples/gcn_15d.py | 355 ++++++++++++++++++++++----------------- 3 files changed, 250 insertions(+), 180 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 65bb4d5..c9bb1d8 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -70,10 +70,12 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand_rows = torch.arange(mb_count * batch_size, 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 = 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])) batches_expand = batches_expand.to_sparse_csr() @@ -84,7 +86,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, total_start_timer.record() for i in range(n_layers): - print(f"Sampling layer {i}", flush=True) + # print(f"Sampling layer {i}", flush=True) if i == 0: nnz = batch_size else: @@ -96,15 +98,19 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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 = 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_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) current_frontier = current_frontier.to_sparse_csr() + print(f"layer {i} current_frontier.size: {current_frontier.size()}", flush=True) # Expand batches matrix if baseline_compare: total_start_timer.record() @@ -142,10 +148,12 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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 = torch.sparse_coo_tensor(next_frontier_idxs, + # next_frontier_values, + # size=(nnz * mb_count, node_count_total)) + # # size=(batch_size * mb_count, node_count_total)) + next_frontier = sparse_coo_tensor_gpu(next_frontier_idxs, next_frontier_values, + torch.Size([nnz * mb_count, node_count_total])) next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) # current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ @@ -166,10 +174,14 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * nnz) adj_sample_cols = adj_sample_cols[nnz_mask] adj_matrices_indices = torch.stack((adj_sample_rows, adj_sample_cols)) - adj_matrices_values = torch.cuda.DoubleTensor(adj_sample_rows.size(0)).fill_(1) - - adj_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_matrices_values = torch.cuda.DoubleTensor(adj_sample_rows.size(0)).fill_(1) + adj_matrices_values = torch.cuda.FloatTensor(adj_sample_rows.size(0)).fill_(1) + + # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, + # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) + adj_matrix_sample = sparse_coo_tensor_gpu(adj_matrices_indices, adj_matrices_values, + torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) + print(f"layer: {i} adj_matrix_sample: {adj_matrix_sample}", flush=True) # adj_matrices[i] = adj_matrix_sample adj_matrices[i] = adj_matrix_sample.to_sparse_csr() frontiers[i + 1] = next_frontier_select.clone() @@ -184,7 +196,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - if timing: + if False and 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": v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index d40e597..86960ff 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -35,7 +35,11 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=Fals timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) def csr_allreduce(mat, left, right, rank, name=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 @@ -54,7 +58,7 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): for req in reqs: req.wait() - torch.cuda.synchronize() + # 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: @@ -72,10 +76,15 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): 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() + # torch.cuda.synchronize() + start_time(start_timer) mat_recv = torch.sparse_csr_tensor(rows_recv, cols_recv, vals_recv, size=mat.size()) + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-csrinst-{name}") + start_time(start_timer) mat = mat + mat_recv + stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-sum-{name}") if rank <= mid_rank: right = mid_rank @@ -570,7 +579,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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, recv_rank, tag=3) + dist.P2POp(dist.isend, matb_send_values.float(), recv_rank, tag=3) send_idx += 4 else: # matb_recv_crows = matb_send_crows.clone() @@ -661,7 +670,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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.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)) @@ -854,8 +864,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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-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}"] = [] @@ -874,7 +884,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) + matc = csr_allreduce(matc, rank_row_start, rank_row_stop, rank, "prob", timing_dict) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-{name}", barrier=True) start_time(start_timer) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 8bd6c3c..d01434b 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -67,29 +67,30 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa self.timings["gemm_w"] = [] self.timings["aggr"] = [] - # input layer - self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) - # hidden layers - for i in range(n_layers - 2): - self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) - # output layer - self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) - # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) - - # self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) - # for _ in range(n_layers - 2): - # self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) - # self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) + # # input layer + # self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # # hidden layers + # for i in range(n_layers - 2): + # self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # # output layer + # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) + # # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + + self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) + for _ in range(n_layers - 2): + self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) + self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): - # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv - # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + graphs[l].t_() + nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv + edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv # edge_index = graphs[l] # SAGEConvfake # edge_index, _, size = graphs[l] # quiver - # h = self.layers[l](h, edge_index) # SAGEConv - h = layer(self, graphs[l], h, epoch) # GCNConv + 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) @@ -138,13 +139,14 @@ def evaluate(self, graph, features, test_idx, labels): torch.FloatTensor(edge_index.size(1)).fill_(1.0), size) adj_batch = adj_batch.t().coalesce() - h_batch = features[n_id] + h = features[n_id] - h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + # 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 + # hs.append(h) # GCNConv + hs.append(h[:batch_size]) # SAGEConv features = torch.cat(hs, dim=0) return features @@ -309,7 +311,6 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep with torch.no_grad(): # Column partitions am_partitions, vtx_indices = split_coo(adj_matrix, node_count, n_per_proc, 1) - torch.cuda.synchronize() # 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) @@ -356,9 +357,10 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep def one5d_partition_mb(rank, size, batches, replication, mb_count): rank_c = rank // replication - # batch_partitions = torch.split(batches, math.ceil(float(mb_count / (size / replication))), dim=0) 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 @@ -479,9 +481,6 @@ def main(args, batches=None): # inputs = data.x.to(device) inputs = data.x data.y = data.y.squeeze().to(device) - print(f"data.y: {data.y}", flush=True) - print(f"data.y.size: {data.y.size()}", flush=True) - print(f"data.y.dtype: {data.y.dtype}", flush=True) # inputs.requires_grad = True # data.y = data.y.to(device) num_features = dataset.num_features @@ -513,110 +512,174 @@ def main(args, batches=None): if rank_row >= (size // args.replication): return - print("start partitioning", flush=True) - torch.cuda.synchronize() if args.dataset == "ogbn-papers100M" or "Protein" in args.dataset: if rank % args.gpu == 0: - inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) - # 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) - - for i in range(1, args.gpu): - print(f"iteration i: {i}", flush=True) - 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_() + if True or args.dataset != "ogbn-papers100M": + inputs = inputs.to(torch.device("cpu")) + adj_matrix = adj_matrix.to(torch.device("cpu")) + # 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) + + for i in range(1, args.gpu): + print(f"iteration i: {i}", flush=True) + 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(data.y.float(), dst=dst_gpu) + # torch.save(g_send_meta, f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_meta.pt") + # torch.save(g_send._indices(), f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_indices.pt") + # torch.save(g_send._values(), f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_values.pt") + # torch.save(features_send, f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}features_loc.pt") + + 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) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t_() + print("normalizing", flush=True) + g_loc = g_loc.to(device) + features_loc = features_loc.to(device) + edge_count = adj_matrix.size(1) + else: + print(f"loading", flush=True) + g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", + map_location=torch.device("cpu")) + g_loc_indices = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_indices.pt", + map_location=torch.device("cpu")) + g_loc_values = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_values.pt", + map_location=torch.device("cpu")) + features_loc = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}features_loc.pt", + map_location=torch.device("cpu")) + + g_loc_meta = g_loc_meta.to(device) + g_loc_indices = g_loc_indices.to(device) + g_loc_values = g_loc_values.to(device) + features_loc = features_loc.to(device) + edge_count = g_loc_meta[3].item() + num_features = g_loc_meta[4].item() + num_classes = g_loc_meta[5].item() + # g_loc = sparse_coo_tensor_gpu(g_loc_indices, g_loc_values, + # torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, + torch.Size([g_loc_meta[1], g_loc_meta[2]])) 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(data.y.float(), dst=dst_gpu) - - 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) - print(f"train_idx_len: {train_idx_len} train_idx: {train_idx}", flush=True) - print(f"test_idx_len: {test_idx_len} test_idx: {test_idx}", flush=True) - dist.send(train_idx, dst=dst_gpu) - dist.send(test_idx, dst=dst_gpu) - # features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ # args.replication, args.normalize) - print("coalescing", flush=True) - g_loc = g_loc.coalesce().t_() - print("normalizing", flush=True) - g_loc = g_loc.to(device) - features_loc = features_loc.to(device) - g_loc = g_loc.double() - edge_count = adj_matrix.size(1) + # 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) - g_loc_meta = torch.cuda.LongTensor(6).fill_(0) - dist.recv(g_loc_meta, src=src_gpu) + if True or args.dataset != "ogbn-papers100M": + 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) - g_loc_indices = torch.cuda.LongTensor(2, g_loc_meta[0].item()) - dist.recv(g_loc_indices, 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) + + # 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_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) - dist.recv(g_loc_values, src=src_gpu) + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + edge_count = g_loc_meta[3].item() - 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) + 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) else: - rows = g_loc_meta[1].item() + (g_loc_meta[2].item() - proc_row * g_loc_meta[1].item()) - print(f"g_loc_meta: {g_loc_meta} rows: {rows}", flush=True) - inputs = torch.cuda.FloatTensor(rows, num_features) - dist.recv(inputs, src=src_gpu) - - # 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, \ - size=torch.Size([g_loc_meta[1], g_loc_meta[2]])) - 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) - print(f"recv train_idx: {train_idx}", flush=True) - print(f"recv test_idx: {test_idx}", flush=True) + print(f"loading", flush=True) + # Load data to CPU instead of GPU bc multiple processes spawn for some reason, + # and pytorch reserves less memory causing OOMs + g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", + map_location=torch.device("cpu")) + g_loc_indices = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_indices.pt", + map_location=torch.device("cpu")) + g_loc_values = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_values.pt", + map_location=torch.device("cpu")) + inputs = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}features_loc.pt", + map_location=torch.device("cpu")) + + g_loc_meta = g_loc_meta.to(device) + g_loc_indices = g_loc_indices.to(device) + g_loc_values = g_loc_values.to(device) + inputs = inputs.to(device) + edge_count = g_loc_meta[3].item() + num_features = g_loc_meta[4].item() + num_classes = g_loc_meta[5].item() + g_loc = sparse_coo_tensor_gpu(g_loc_indices, g_loc_values, + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + data.y = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/datay.pt") + train_idx = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/train_idx.pt") + test_idx = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/test_idx.pt") + data.y = data.y.to(device) + train_idx = train_idx.to(device) + test_idx = test_idx.to(device) torch.cuda.synchronize() features_loc = inputs.to(device) else: @@ -674,7 +737,7 @@ def main(args, batches=None): # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model - torch.manual_seed(0) + # torch.manual_seed(0) model = GCN(num_features, args.n_hidden, num_classes, @@ -690,7 +753,7 @@ def main(args, batches=None): model = model.to(device) # use optimizer - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) dur = [] # torch.manual_seed(0) @@ -706,12 +769,13 @@ def main(args, batches=None): print(f"rank_n_bulkmb: {rank_n_bulkmb}") # g_loc = g_loc.cpu() + g_loc = g_loc.to_sparse_csr() 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') + # 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): @@ -737,27 +801,6 @@ def main(args, batches=None): batches_all = train_nid[vertex_perm] torch.cuda.nvtx.range_pop() # construct-batches - torch.cuda.nvtx.range_push("nvtx-extract-batches") - last_batch = False - 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 - batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) - batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) - - batches_indices_rows = torch.arange(batches_loc.size(0)).to(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)) - if epoch >= 1: - epoch_start = time.time() - torch.cuda.nvtx.range_pop() # extract-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) @@ -785,7 +828,7 @@ def main(args, batches=None): 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.DoubleTensor(batches_loc.size(1) * batches_loc.size(0)).fill_(1.0) + batches_values = torch.cuda.DoubleTensor(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) @@ -847,21 +890,24 @@ def main(args, batches=None): batch_select_min = i * args.batch_size batch_select_max = (i + 1) * args.batch_size - rank_n_bulkmb_row = rank_n_bulkmb // args.replication - if rank_col == args.replication - 1: - rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) + # rank_n_bulkmb_row = rank_n_bulkmb // args.replication + # if rank_col == args.replication - 1: + # rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) - batch_select_min += batch_select_size * rank_col * rank_n_bulkmb_row - batch_select_max += batch_select_size * rank_col * rank_n_bulkmb_row + batch_select_min += batch_select_size * rank_col * rank_n_bulkmb + batch_select_max += batch_select_size * rank_col * rank_n_bulkmb batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) + print(f"rank_n_bulkmb: {rank_n_bulkmb}", flush=True) + print(f"i: {i} batch_vtxs: {batch_vtxs}", flush=True) + print(f"i: {i} batch_select_min: {batch_select_min} max: {batch_select_max}", flush=True) # src_vtxs = frontiers[-1][i].view(-1) src_select_size = args.batch_size * (args.samp_num ** (args.n_layers - 1)) src_select_min = i * args.batch_size * (args.samp_num ** (args.n_layers - 1)) src_select_max = (i + 1) * args.batch_size * (args.samp_num ** (args.n_layers - 1)) - src_select_min += src_select_size * rank_col * rank_n_bulkmb_row - src_select_max += src_select_size * rank_col * rank_n_bulkmb_row + src_select_min += src_select_size * rank_col * rank_n_bulkmb + src_select_max += src_select_size * rank_col * rank_n_bulkmb src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) # print(f"src_vtxs: {src_vtxs} src_vtxs.size: {src_vtxs.size()} zero_count: {(src_vtxs == 0).sum()}", flush=True) @@ -870,13 +916,15 @@ def main(args, batches=None): adj_select_size = args.batch_size * (args.samp_num ** l) 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 += adj_select_size * rank_col * rank_n_bulkmb_row - adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb_row + adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb + adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb adj_row_select_max = min(adj_row_select_max, adj_matrices_bulk[l].size(0)) if epoch >= 1: start_time(start_inner_timer) + print(f"l: {l} adj_matrices_bulk.size: {adj_matrices_bulk[l].size()}", flush=True) + print(f"l: {l} adj_row_select_min: {adj_row_select_min} max: {adj_row_select_max}", flush=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] @@ -1017,9 +1065,8 @@ def main(args, batches=None): if epoch >= 1: start_time(start_inner_timer) torch.cuda.nvtx.range_push("nvtx-loss") - print(f"logits: {logits[:args.batch_size]}", flush=True) - 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, data.y[batch_vtxs].long()) # GCNConv + loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs].long()) # SAGEConv # loss = F.nll_loss(logits[:args.batch_size], data.y[seeds].long()) # quiver torch.cuda.nvtx.range_pop() # nvtx-loss @@ -1093,7 +1140,7 @@ def main(args, batches=None): # print(f"adj1_nnz.median: {np.median(adj1_nnzs)}") # print(f"adj_nnzs.median: {np.median(adj_nnzs)}") print(f"total: {np.sum(dur)}", flush=True) - if args.dataset != "Amazon" and ("Protein" not in args.dataset): + if args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): model = model.cpu() train_nid = train_nid.cpu() test_nid = test_nid.cpu() From cf6d12d92b4f9c1df9aafa0eb5343f4529f91ac9 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 4 Oct 2023 01:54:52 -0700 Subject: [PATCH 152/178] Optimized sparse allreduce for sage --- cagnet/samplers/sage.py | 6 +- cagnet/samplers/utils.py | 25 +- examples/gcn_15d.py | 21 +- sparse-extension/setup.py | 2 +- sparse-extension/sparse_coo_tensor.cu | 2725 +++++++++++++------------ 5 files changed, 1441 insertions(+), 1338 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index c9bb1d8..85d6554 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -86,7 +86,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, total_start_timer.record() for i in range(n_layers): - # print(f"Sampling layer {i}", flush=True) + print(f"Sampling layer {i}", flush=True) if i == 0: nnz = batch_size else: @@ -110,7 +110,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, nnz = batch_size * (frontier_size ** i) current_frontier = current_frontier.to_sparse_csr() - print(f"layer {i} current_frontier.size: {current_frontier.size()}", flush=True) # Expand batches matrix if baseline_compare: total_start_timer.record() @@ -181,7 +180,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) adj_matrix_sample = sparse_coo_tensor_gpu(adj_matrices_indices, adj_matrices_values, torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) - print(f"layer: {i} adj_matrix_sample: {adj_matrix_sample}", flush=True) # adj_matrices[i] = adj_matrix_sample adj_matrices[i] = adj_matrix_sample.to_sparse_csr() frontiers[i + 1] = next_frontier_select.clone() @@ -196,7 +194,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.cuda.synchronize() total_time = total_start_timer.elapsed_time(total_stop_timer) print(f"total_time: {total_time}", flush=True) - if False and timing: + 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": v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 86960ff..3170fe5 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -10,7 +10,8 @@ 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 + sparse_csr_tensor_gpu, nsparse_spgemm, rearrange_rows_gpu, \ + reduce_sum_gpu timing = True @@ -80,10 +81,28 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): # torch.cuda.synchronize() start_time(start_timer) - mat_recv = torch.sparse_csr_tensor(rows_recv, cols_recv, vals_recv, size=mat.size()) + 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) - mat = mat + mat_recv + # mat = mat + mat_recv + + # 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 = torch.sparse_csr_tensor(mat_sum_crows, mat_sum_cols, mat_sum_vals, mat.size()) stop_time_add(start_timer, stop_timer, timing_dict, f"spgemm-reduce-sum-{name}") if rank <= mid_rank: diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index d01434b..934bfc5 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -851,7 +851,7 @@ def main(args, batches=None): if epoch >= 1: start_time(start_timer) - # print("Sampling", flush=True) + print("Sampling", flush=True) torch.cuda.nvtx.range_push("nvtx-sampling") if args.sample_method == "ladies": current_frontier, next_frontier, adj_matrices_bulk = \ @@ -878,7 +878,7 @@ def main(args, batches=None): start_time(start_timer) # g_loc = g_loc.coalesce() - # print("Training", flush=True) + print("Training", flush=True) torch.cuda.nvtx.range_push("nvtx-training") for i in range(rank_n_bulkmb): @@ -894,20 +894,17 @@ def main(args, batches=None): # if rank_col == args.replication - 1: # rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) - batch_select_min += batch_select_size * rank_col * rank_n_bulkmb - batch_select_max += batch_select_size * rank_col * rank_n_bulkmb + # batch_select_min += batch_select_size * rank_col * rank_n_bulkmb + # batch_select_max += batch_select_size * rank_col * rank_n_bulkmb batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) - print(f"rank_n_bulkmb: {rank_n_bulkmb}", flush=True) - print(f"i: {i} batch_vtxs: {batch_vtxs}", flush=True) - print(f"i: {i} batch_select_min: {batch_select_min} max: {batch_select_max}", flush=True) # src_vtxs = frontiers[-1][i].view(-1) src_select_size = args.batch_size * (args.samp_num ** (args.n_layers - 1)) src_select_min = i * args.batch_size * (args.samp_num ** (args.n_layers - 1)) src_select_max = (i + 1) * args.batch_size * (args.samp_num ** (args.n_layers - 1)) - src_select_min += src_select_size * rank_col * rank_n_bulkmb - src_select_max += src_select_size * rank_col * rank_n_bulkmb + # src_select_min += src_select_size * rank_col * rank_n_bulkmb + # src_select_max += src_select_size * rank_col * rank_n_bulkmb src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) # print(f"src_vtxs: {src_vtxs} src_vtxs.size: {src_vtxs.size()} zero_count: {(src_vtxs == 0).sum()}", flush=True) @@ -916,15 +913,13 @@ def main(args, batches=None): adj_select_size = args.batch_size * (args.samp_num ** l) 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 += adj_select_size * rank_col * rank_n_bulkmb - adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb + # adj_row_select_min += adj_select_size * rank_col * rank_n_bulkmb + # adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb adj_row_select_max = min(adj_row_select_max, adj_matrices_bulk[l].size(0)) if epoch >= 1: start_time(start_inner_timer) - print(f"l: {l} adj_matrices_bulk.size: {adj_matrices_bulk[l].size()}", flush=True) - print(f"l: {l} adj_row_select_min: {adj_row_select_min} max: {adj_row_select_max}", flush=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] diff --git a/sparse-extension/setup.py b/sparse-extension/setup.py index dd35dac..d335145 100644 --- a/sparse-extension/setup.py +++ b/sparse-extension/setup.py @@ -3,5 +3,5 @@ setup(name='sparse_coo_tensor_cpp', 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"])], + 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.cu b/sparse-extension/sparse_coo_tensor.cu index 8badd94..d20f28f 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -62,1274 +62,1274 @@ using namespace at::sparse; #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); - } - - 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); - } - - 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); - } - - __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) { - // printf("infinite loop key %d adr %d offset: %d hash: %d d_check[adr] %d bcol: %d max_row_nz: %d\n", key, adr, offset, hash, d_check[adr], bcol, max_row_nz); - 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); - } - - 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])); - // torch::Tensor d_fail_count_tens = torch::zeros({1}, torch::TensorOptions() - // .dtype(torch::kInt32) - // .device(torch::kCUDA)); - // torch::Tensor d_fail_perm_tens = torch::zeros({bin->bin_size[i]}, torch::TensorOptions() - // .dtype(torch::kInt32) - // .device(torch::kCUDA)); - // int *d_fail_count = d_fail_count_tens.data(); - // int *d_fail_perm = d_fail_perm_tens.data(); - 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)); - - // torch::Tensor d_check_tens = torch::zeros({(long)table_size}, torch::TensorOptions() - // .dtype(torch::kInt32) - // .device(torch::kCUDA)); - // int *d_check = d_check_tens.data(); - 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)); - - torch::Tensor d_check_tens = torch::zeros({table_size}, torch::TensorOptions() - .dtype(torch::kInt32) - .device(torch::kCUDA)); - torch::Tensor d_value_tens = torch::zeros({table_size}, torch::TensorOptions() - .dtype(torch::kFloat32) - .device(torch::kCUDA)); - int *d_check = d_check_tens.data(); - float *d_value = d_value_tens.data(); // assumes real type is float - - 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(); -} +// 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); +// } +// +// 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); +// } +// +// 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); +// } +// +// __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) { +// // printf("infinite loop key %d adr %d offset: %d hash: %d d_check[adr] %d bcol: %d max_row_nz: %d\n", key, adr, offset, hash, d_check[adr], bcol, max_row_nz); +// 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); +// } +// +// 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])); +// // torch::Tensor d_fail_count_tens = torch::zeros({1}, torch::TensorOptions() +// // .dtype(torch::kInt32) +// // .device(torch::kCUDA)); +// // torch::Tensor d_fail_perm_tens = torch::zeros({bin->bin_size[i]}, torch::TensorOptions() +// // .dtype(torch::kInt32) +// // .device(torch::kCUDA)); +// // int *d_fail_count = d_fail_count_tens.data(); +// // int *d_fail_perm = d_fail_perm_tens.data(); +// 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)); +// +// // torch::Tensor d_check_tens = torch::zeros({(long)table_size}, torch::TensorOptions() +// // .dtype(torch::kInt32) +// // .device(torch::kCUDA)); +// // int *d_check = d_check_tens.data(); +// 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)); +// +// torch::Tensor d_check_tens = torch::zeros({table_size}, torch::TensorOptions() +// .dtype(torch::kInt32) +// .device(torch::kCUDA)); +// torch::Tensor d_value_tens = torch::zeros({table_size}, torch::TensorOptions() +// .dtype(torch::kFloat32) +// .device(torch::kCUDA)); +// int *d_check = d_check_tens.data(); +// float *d_value = d_value_tens.data(); // assumes real type is float +// +// 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_searchf(double *arr, double val, int imin, int imax) { @@ -1768,68 +1768,68 @@ std::vector nsparse_spgemm( int32_t k, int32_t m) { - sfBIN bin; + // 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(); + // /* 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(); + // 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); + // init_bin(&bin, n); + // // 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(); - // 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); + // set_max_bin(d_a_crow_indices, d_a_col_indices, d_b_crow_indices, &bin, n); - 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); - - /* Set bin */ - set_min_bin(&bin, n); + // 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); + + // /* Set bin */ + // set_min_bin(&bin, n); - // 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)); + // // 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(); + // 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); + // /* 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); - release_bin(bin); + // release_bin(bin); - CHECK_ERROR("nsparse_spgemm") + // CHECK_ERROR("nsparse_spgemm") - return {C_crow_indices, C_col_indices, C_values}; - // return {c_csr_offsets, c_columns, c_values}; + // return {C_crow_indices, C_col_indices, C_values}; + // // return {c_csr_offsets, c_columns, c_values}; } std::vector spgemm_gpu( @@ -2629,6 +2629,96 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols 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 *matb_crows, + long *matc_cols, long *mata_cols, long *matb_cols, + long *mata_nnz_rows, long *matb_nnz_rows, + int mata_nnz_row_count, int matb_nnz_row_count) { + + int id = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (int i = id; i < mata_nnz_row_count; i += stride) { + int row_id = mata_nnz_rows[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 < matb_nnz_row_count; i += stride) { + int row_id = matb_nnz_rows[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 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(), + matb_crows.data(), + matc_cols.data(), + mata_cols.data(), + matb_cols.data(), + mata_nnz_rows.data(), + matb_nnz_rows.data(), + mata_nnz_row_count, + matb_nnz_row_count); + + // ReduceSum<<>>(matc_crows.data(), + // matb_crows.data(), + // matc_cols.data(), + // matb_cols.data(), + // ROW_COUNT); + + CHECK_ERROR("reduce 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"); @@ -2653,4 +2743,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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("reduce_sum_gpu", &reduce_sum_gpu, "Sum for SAGE sa-spgemm reduction"); } From f16f436101e313e167a4ab29e5d68527fd77e96e Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 5 Oct 2023 01:41:00 -0700 Subject: [PATCH 153/178] Split sampling select within process row + ring allgather --- cagnet/samplers/sage.py | 7 +- cagnet/samplers/utils.py | 392 +++++++++++++++++++++++--- examples/gcn_15d.py | 7 +- sparse-extension/sparse_coo_tensor.cu | 49 ++-- 4 files changed, 384 insertions(+), 71 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 85d6554..1835b9d 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -157,6 +157,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) # current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ # current_frontier.values().bool()).view(current_frontier._nnz(), 1) + timing_dict["frontier-row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) + + start_time(start_timer) current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) if i == 0: frontiers[0] = current_frontier_select.clone() @@ -184,7 +187,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrices[i] = adj_matrix_sample.to_sparse_csr() frontiers[i + 1] = next_frontier_select.clone() current_frontier = next_frontier - timing_dict["row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) + timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) del next_frontier_select del current_frontier_select @@ -196,7 +199,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, print(f"total_time: {total_time}", 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": + 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"): v_tens = torch.cuda.FloatTensor(1).fill_(sum(v)) v_tens_recv = [] for i in range(size): diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 3170fe5..52f4fab 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -35,6 +35,38 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=Fals if timing_dict is not None: timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) +def ring_allgather(input_tensor, rank, proc_list): + output_tensor_recv = [None] * len(proc_list) + + rank_idx = proc_list.index(rank) + current_buffer = input_tensor + for i in range(len(proc_list)): + recv_rank = proc_list[rank_idx - 1] + send_rank = proc_list[(rank_idx + 1) % len(proc_list)] + + # Exchange sizes + ops = [None] * 2 + recv_size = torch.cuda.IntTensor(1).fill_(0) + send_size = torch.cuda.IntTensor(1).fill_(current_buffer.size(0)) + ops[0] = dist.P2POp(dist.isend, send_size, send_rank) + ops[1] = dist.P2POp(dist.irecv, recv_size, recv_rank) + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + + # Alloc buffer + recv_buffer = torch.cuda.IntTensor(recv_size.item()) + ops[0] = dist.P2POp(dist.isend, current_buffer, send_rank) + ops[1] = dist.P2POp(dist.irecv, recv_buffer, recv_rank) + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + + current_buffer = recv_buffer + output_tensor_recv[rank_idx - 1 - i] = recv_buffer + + return output_tensor_recv + def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -1072,10 +1104,10 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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 + 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 next_frontier = torch.sparse_coo_tensor(indices=p._indices(), values=torch.cuda.LongTensor(p._nnz()).fill_(0), @@ -1094,7 +1126,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, zero_idxs = p._indices()[0, (p._values() == 0)].unique() frontier_nnz_sizes[zero_idxs] = 0 - timing_dict["pre-loop"].append(stop_time(start_timer, stop_timer)) + 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 @@ -1124,26 +1156,26 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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["prob-rowsum"].append(stop_time(start_timer, stop_timer)) + timing_dict["sample-prob-rowsum"].append(stop_time(start_timer, stop_timer)) 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_() - timing_dict["gen-darts"].append(stop_time(start_timer, stop_timer)) + 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)) - timing_dict["dart-throw"].append(stop_time(start_timer, stop_timer)) + 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.size(0), p._nnz()) # n_darts * mb_count, p._nnz()) - timing_dict["filter-darts"].append(stop_time(start_timer, stop_timer)) + timing_dict["sample-filter-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) @@ -1155,7 +1187,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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["add-to-frontier"].append(stop_time(start_timer, stop_timer)) + 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() @@ -1166,22 +1198,38 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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) - timing_dict["count-samples"].append(stop_time(start_timer, stop_timer)) + timing_dict["sample-count-samples"].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-preproc"].append(stop_time(start_timer, stop_timer)) + timing_dict["sample-select-preproc"].append(stop_time(start_timer, stop_timer)) start_time(select_start_timer) - if rank_col == 0 and overflowed_minibatches: + rank_n_bulkmb = int(next_frontier.size(0) / replication) + mb_start = rank_col * rank_n_bulkmb + mb_stop = min((rank_col + 1) * rank_n_bulkmb, next_frontier.size(0)) + rank_n_bulkmb = mb_stop - mb_start + if rank_col == replication - 1: + rank_n_bulkmb = next_frontier.size(0) - rank_n_bulkmb * (replication - 1) + mb_stop = next_frontier.size(0) + select_nnz_mask = (next_frontier._indices()[0,:] >= mb_start) & \ + (next_frontier._indices()[0,:] < mb_stop) + # if rank_col == 0 and overflowed_minibatches: + overflow_select = overflow[mb_start:mb_stop] + sampled_count = sampled_count[mb_start:mb_stop] + overflowed_minibatches = (overflow_select > 0).any() + dart_hits_count_select = dart_hits_count[select_nnz_mask] + next_frontier_select_idxs = next_frontier._indices()[:,select_nnz_mask] + next_frontier_select_idxs[0,:] -= mb_start + if 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) + # ps_overflow = torch.cumsum(overflow, dim=0) + ps_overflow = torch.cumsum(overflow_select, dim=0) total_overflow = ps_overflow[-1].item() # n_darts_select = total_overflow // replication # if rank_col == replication - 1: @@ -1190,23 +1238,24 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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_count.reciprocal().double() + dart_hits_inv = dart_hits_count_select.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 = torch.sparse_coo_tensor(indices=next_frontier._indices(), - # values=dart_hits_inv, - # size=(p.size(0), node_count_total)) - dart_hits_inv_mtx = sparse_coo_tensor_gpu(next_frontier._indices(), dart_hits_inv, - torch.Size([p.size(0), node_count_total])) + # dart_hits_inv_mtx = sparse_coo_tensor_gpu(next_frontier._indices(), dart_hits_inv, + # torch.Size([p.size(0), node_count_total])) + dart_hits_inv_mtx = sparse_coo_tensor_gpu(next_frontier_select_idxs, dart_hits_inv, + torch.Size([rank_n_bulkmb, 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_sum.scatter_add_(0, next_frontier._indices()[0, :], dart_hits_inv) + # dart_hits_inv_sum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) + dart_hits_inv_sum = torch.cuda.DoubleTensor(rank_n_bulkmb).fill_(0) dart_hits_inv_nnzidxs = dart_hits_inv.nonzero().squeeze() - indices_nnz_idxs = next_frontier._indices()[0,dart_hits_inv_nnzidxs] + # indices_nnz_idxs = next_frontier._indices()[0,dart_hits_inv_nnzidxs] + indices_nnz_idxs = next_frontier_select_idxs[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)) @@ -1220,7 +1269,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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) + rank_n_bulkmb, n_darts_select) + # p.size(0), n_darts_select) # mb_count, total_overflow) timing_dict["select-computedarts"].append(stop_time(start_timer, stop_timer)) @@ -1228,23 +1278,30 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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()) + # throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count, n_darts_select, + # next_frontier._nnz()) + throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count_select, n_darts_select, + next_frontier_select_idxs.size(1)) 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_values = torch.logical_or(dart_hits_count, next_frontier._values()).int() + next_frontier_values = torch.logical_or(dart_hits_count_select, + next_frontier._values()[select_nnz_mask]).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])) + # next_frontier_tmp = sparse_coo_tensor_gpu(next_frontier._indices(), next_frontier_values, + # torch.Size([p.size(0), node_count_total])) + next_frontier_tmp = sparse_coo_tensor_gpu(next_frontier_select_idxs, next_frontier_values, + torch.Size([rank_n_bulkmb, 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 = torch.cuda.IntTensor(p.size(0)).fill_(0) + sampled_count = torch.cuda.IntTensor(rank_n_bulkmb).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] @@ -1253,8 +1310,10 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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() + # overflow = torch.clamp(sampled_count - frontier_size, min=0).int() + overflow_select = torch.clamp(sampled_count - frontier_size, min=0).int() + # overflowed_minibatches = (overflow > 0).any() + overflowed_minibatches = (overflow_select > 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)) @@ -1273,10 +1332,266 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, dist.barrier(group=row_groups[rank_c]) - timing_dict["dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) + timing_dict["sample-dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) + start_time(start_timer) # 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]) + # dist.broadcast(dart_hits_count, src=rank_c * replication, group=row_groups[rank_c]) + # max_len = 0 + # max_len_samples = 0 + # for i in range(replication): + # mb_start_i = i * int(next_frontier.size(0) / replication) + # mb_stop_i = (i + 1) * int(next_frontier.size(0) / replication) + # if i == replication - 1: + # mb_stop_i = next_frontier.size(0) + # i_nnz_mask = (next_frontier._indices()[0,:] >= mb_start_i) & \ + # (next_frontier._indices()[0,:] < mb_stop_i) + # max_len = max(i_nnz_mask.size(0), max_len) + # max_len_samples = max(mb_stop_i - mb_start_i, max_len_samples) + + # dart_hits_count_recv = torch.cuda.IntTensor(max_len * replication) + # sampled_count_recv = torch.cuda.IntTensor(max_len_samples * replication) + + # if dart_hits_count_select.size(0) < max_len: + # padding = torch.cuda.IntTensor(max_len - dart_hits_count_select.size(0)).fill_(-1) + # print(f"dart.size: {dart_hits_count_select.size()} dartpadding.size: {padding.size()}", flush=True) + # dart_hits_count_select = torch.cat((dart_hits_count_select, padding), dim=0) + # if sampled_count.size(0) < max_len_samples: + # padding = torch.cuda.IntTensor(max_len_samples - sampled_count.size(0)).fill_(-1) + # print(f"sample.size: {sampled_count.size()} samplepadding.size: {padding.size()}", flush=True) + # sampled_count = torch.cat((sampled_count, padding), dim=0) + + dart_lens = [] + sample_lens = [] + for i in range(replication): + mb_start_i = i * int(next_frontier.size(0) / replication) + mb_stop_i = (i + 1) * int(next_frontier.size(0) / replication) + if i == replication - 1: + mb_stop_i = next_frontier.size(0) + i_nnz_mask = (next_frontier._indices()[0,:] >= mb_start_i) & \ + (next_frontier._indices()[0,:] < mb_stop_i) + # max_len = max(i_nnz_mask.size(0), max_len) + dart_lens.append(i_nnz_mask.size(0)) + # max_len_samples = max(mb_stop_i - mb_start_i, max_len_samples) + sample_lens.append(mb_stop_i - mb_start_i) + + dart_hits_count_recv = [] + sample_count_recv = [] + for i in range(replication): + dart_hits_count_recv.append(torch.cuda.IntTensor(dart_lens[i])) + sample_count_recv.append(torch.cuda.IntTensor(sample_lens[i])) + timing_dict["sample-gather-counts-preproc"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # dist.all_gather_into_tensor(dart_hits_count_recv, dart_hits_count_select, group=row_groups[rank_c]) + # dist.all_gather_into_tensor(sampled_count_recv, sampled_count, group=row_groups[rank_c]) + # dart_hits_count = dart_hits_count_recv[(dart_hits_count_recv != -1)] + # sampled_count = sampled_count_recv[(sampled_count_recv != -1)] + row_procs = list(range(rank_c * replication, (rank_c + 1) * replication)) + dart_hits_count_recv = ring_allgather(dart_hits_count_select, rank, row_procs) + sampled_count_recv = ring_allgather(sampled_count, rank, row_procs) + dart_hits_count = torch.cat(dart_hits_count_recv, dim=0) + sampled_count = torch.cat(sampled_count_recv, dim=0) + 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,:]) + p._values()[filled_minibatches_mask] = 0 + timing_dict["sample-set-probs"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() + next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] + 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_nnzidxs, \ + next_frontier._values()[next_frontier_nnzvals].int()) + 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)) + + # print(f"iter_count: {iter_count}") + # print(f"selection_iter_count: {selection_iter_count}") + + 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])) + sa_masks.fill_(0) + row_select_mtx = row_select_mtx.to_sparse_csr() + # adj_matrix = adj_matrix.to_sparse_csr() + print(f"adj_matrix.layout: {adj_matrix.layout}") + 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 + timing_dict["sample-gather-counts-preproc"].append(stop_time(start_timer, stop_timer)) + + start_time(start_timer) + # dist.all_gather_into_tensor(dart_hits_count_recv, dart_hits_count_select, group=row_groups[rank_c]) + # dist.all_gather_into_tensor(sampled_count_recv, sampled_count, group=row_groups[rank_c]) + # dart_hits_count = dart_hits_count_recv[(dart_hits_count_recv != -1)] + # sampled_count = sampled_count_recv[(sampled_count_recv != -1)] + dist.all_gather_into_tensor(dart_hits_count, dart_hits_count_select, group=row_groups[rank_c]) + dist.all_gather_into_tensor(sampled_count, sampled_count, group=row_groups[rank_c]) next_frontier_values = torch.logical_or( dart_hits_count, next_frontier._values()).int() @@ -1285,6 +1600,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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 @@ -1293,7 +1609,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, filled_minibatches = sampled_count == frontier_size filled_minibatches_mask = torch.gather(filled_minibatches, 0, p._indices()[0,:]) p._values()[filled_minibatches_mask] = 0 - timing_dict["set-probs"].append(stop_time(start_timer, stop_timer)) + timing_dict["sample-set-probs"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() @@ -1303,7 +1619,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, sampled_count.scatter_add_(0, next_frontier_nnzidxs, \ next_frontier._values()[next_frontier_nnzvals].int()) underfull_minibatches = (sampled_count < frontier_size).any() - timing_dict["compute-bool"].append(stop_time(start_timer, stop_timer)) + timing_dict["sample-compute-bool"].append(stop_time(start_timer, stop_timer)) overflow = torch.clamp(sampled_count - frontier_size, min=0).int() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 934bfc5..0871195 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -27,7 +27,6 @@ import torch.nn.functional as F from sparse_coo_tensor_cpp import sort_dst_proc_gpu -import quiver import socket @@ -1125,8 +1124,8 @@ def main(args, batches=None): 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)} extract-select: {np.sum(extract_select_dur)} extract-inst: {np.sum(extract_inst_dur)} extract-coalesce: {np.sum(extract_coalesce_dur)} train: {np.sum(train_dur)}", flush=True) - print(f"feats: {np.sum(selectfeats_dur)} fwd: {np.sum(fwd_dur)} bwd: {np.sum(bwd_dur)}") + 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)}") @@ -1134,7 +1133,7 @@ def main(args, batches=None): # 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)}", flush=True) + print(f"total: {np.sum(dur) / epoch}", flush=True) if args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): model = model.cpu() train_nid = train_nid.cpu() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index d20f28f..d580505 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2657,29 +2657,28 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols // } // } -__global__ void ReduceSum(long *matc_crows, long *mata_crows, long *matb_crows, - long *matc_cols, long *mata_cols, long *matb_cols, - long *mata_nnz_rows, long *matb_nnz_rows, - int mata_nnz_row_count, int matb_nnz_row_count) { +__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 < mata_nnz_row_count; i += stride) { - int row_id = mata_nnz_rows[i]; + 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 < matb_nnz_row_count; i += stride) { - int row_id = matb_nnz_rows[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]; - } - } + // 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, @@ -2689,9 +2688,10 @@ void reduce_sum_gpu(const at::Tensor& matc_crows, const at::Tensor& mata_crows, 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 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); @@ -2701,20 +2701,15 @@ void reduce_sum_gpu(const at::Tensor& matc_crows, const at::Tensor& mata_crows, ReduceSum<<>>(matc_crows.data(), mata_crows.data(), - matb_crows.data(), matc_cols.data(), mata_cols.data(), + ROW_COUNT); + + ReduceSum<<>>(matc_crows.data(), + matb_crows.data(), + matc_cols.data(), matb_cols.data(), - mata_nnz_rows.data(), - matb_nnz_rows.data(), - mata_nnz_row_count, - matb_nnz_row_count); - - // ReduceSum<<>>(matc_crows.data(), - // matb_crows.data(), - // matc_cols.data(), - // matb_cols.data(), - // ROW_COUNT); + ROW_COUNT); CHECK_ERROR("reduce sum error") } From c1f73949f2482b48990b5afec547f53bb18a7148 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 5 Oct 2023 03:27:08 -0700 Subject: [PATCH 154/178] Undid splitting sampling within process row --- cagnet/samplers/sage.py | 16 ++-- cagnet/samplers/utils.py | 157 +++++---------------------------------- 2 files changed, 28 insertions(+), 145 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 1835b9d..ff5a740 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -162,17 +162,18 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, start_time(start_timer) current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) if i == 0: - frontiers[0] = current_frontier_select.clone() + # frontiers[0] = current_frontier_select.clone() + frontiers[0] = current_frontier_select batch_vals = torch.cuda.LongTensor(current_frontier_select.size()).fill_(1) next_frontier_select_vals = next_frontier._values().view(-1) - batch_rows = torch.arange(mb_count * nnz).cuda().view(mb_count * nnz, 1) + batch_rows = torch.arange(mb_count * nnz, device=gpu).view(mb_count * nnz, 1) next_frontier_select_rows = next_frontier._indices()[0,:].view(-1) nnz_mask = next_frontier_select_vals.nonzero().squeeze() adj_sample_rows = next_frontier_select_rows[nnz_mask] - adj_sample_cols = torch.arange(next_frontier_select.numel()).cuda() + adj_sample_cols = torch.arange(next_frontier_select.numel(), device=gpu) adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * nnz) adj_sample_cols = adj_sample_cols[nnz_mask] adj_matrices_indices = torch.stack((adj_sample_rows, adj_sample_cols)) @@ -184,12 +185,13 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, adj_matrix_sample = sparse_coo_tensor_gpu(adj_matrices_indices, adj_matrices_values, torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) # adj_matrices[i] = adj_matrix_sample - adj_matrices[i] = adj_matrix_sample.to_sparse_csr() - frontiers[i + 1] = next_frontier_select.clone() + # frontiers[i + 1] = next_frontier_select.clone() + frontiers[i + 1] = next_frontier_select current_frontier = next_frontier + adj_matrices[i] = adj_matrix_sample.to_sparse_csr() timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) - del next_frontier_select - del current_frontier_select + # del next_frontier_select + # del current_frontier_select # print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) if baseline_compare: diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 52f4fab..b8bd74a 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -35,38 +35,6 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=Fals if timing_dict is not None: timing_dict[range_name].append(stop_time(start_timer, stop_timer, barrier)) -def ring_allgather(input_tensor, rank, proc_list): - output_tensor_recv = [None] * len(proc_list) - - rank_idx = proc_list.index(rank) - current_buffer = input_tensor - for i in range(len(proc_list)): - recv_rank = proc_list[rank_idx - 1] - send_rank = proc_list[(rank_idx + 1) % len(proc_list)] - - # Exchange sizes - ops = [None] * 2 - recv_size = torch.cuda.IntTensor(1).fill_(0) - send_size = torch.cuda.IntTensor(1).fill_(current_buffer.size(0)) - ops[0] = dist.P2POp(dist.isend, send_size, send_rank) - ops[1] = dist.P2POp(dist.irecv, recv_size, recv_rank) - reqs = dist.batch_isend_irecv(ops) - for req in reqs: - req.wait() - - # Alloc buffer - recv_buffer = torch.cuda.IntTensor(recv_size.item()) - ops[0] = dist.P2POp(dist.isend, current_buffer, send_rank) - ops[1] = dist.P2POp(dist.irecv, recv_buffer, recv_rank) - reqs = dist.batch_isend_irecv(ops) - for req in reqs: - req.wait() - - current_buffer = recv_buffer - output_tensor_recv[rank_idx - 1 - i] = recv_buffer - - return output_tensor_recv - def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): start_timer = torch.cuda.Event(enable_timing=True) stop_timer = torch.cuda.Event(enable_timing=True) @@ -1202,34 +1170,18 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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) - rank_n_bulkmb = int(next_frontier.size(0) / replication) - mb_start = rank_col * rank_n_bulkmb - mb_stop = min((rank_col + 1) * rank_n_bulkmb, next_frontier.size(0)) - rank_n_bulkmb = mb_stop - mb_start - if rank_col == replication - 1: - rank_n_bulkmb = next_frontier.size(0) - rank_n_bulkmb * (replication - 1) - mb_stop = next_frontier.size(0) - select_nnz_mask = (next_frontier._indices()[0,:] >= mb_start) & \ - (next_frontier._indices()[0,:] < mb_stop) - # if rank_col == 0 and overflowed_minibatches: - overflow_select = overflow[mb_start:mb_stop] - sampled_count = sampled_count[mb_start:mb_stop] - overflowed_minibatches = (overflow_select > 0).any() - dart_hits_count_select = dart_hits_count[select_nnz_mask] - next_frontier_select_idxs = next_frontier._indices()[:,select_nnz_mask] - next_frontier_select_idxs[0,:] -= mb_start - if overflowed_minibatches: + 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) - ps_overflow = torch.cumsum(overflow_select, dim=0) + ps_overflow = torch.cumsum(overflow, dim=0) total_overflow = ps_overflow[-1].item() # n_darts_select = total_overflow // replication # if rank_col == replication - 1: @@ -1238,24 +1190,19 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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_count_select.reciprocal().double() + 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])) - dart_hits_inv_mtx = sparse_coo_tensor_gpu(next_frontier_select_idxs, dart_hits_inv, - torch.Size([rank_n_bulkmb, node_count_total])) + 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_sum = torch.cuda.DoubleTensor(rank_n_bulkmb).fill_(0) + 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] - indices_nnz_idxs = next_frontier_select_idxs[0,dart_hits_inv_nnzidxs] + 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)) @@ -1269,39 +1216,30 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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, - rank_n_bulkmb, n_darts_select) - # p.size(0), n_darts_select) - # mb_count, total_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()) - throw_darts_select_gpu(dart_select, ps_dart_hits_inv, dart_hits_count_select, n_darts_select, - next_frontier_select_idxs.size(1)) + 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_values = torch.logical_or(dart_hits_count_select, - next_frontier._values()[select_nnz_mask]).int() + 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])) - next_frontier_tmp = sparse_coo_tensor_gpu(next_frontier_select_idxs, next_frontier_values, - torch.Size([rank_n_bulkmb, 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.cuda.IntTensor(rank_n_bulkmb).fill_(0) - # sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) + 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] @@ -1310,10 +1248,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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() - overflow_select = torch.clamp(sampled_count - frontier_size, min=0).int() - # overflowed_minibatches = (overflow > 0).any() - overflowed_minibatches = (overflow_select > 0).any() + 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)) @@ -1336,62 +1272,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(start_timer) # 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]) - # max_len = 0 - # max_len_samples = 0 - # for i in range(replication): - # mb_start_i = i * int(next_frontier.size(0) / replication) - # mb_stop_i = (i + 1) * int(next_frontier.size(0) / replication) - # if i == replication - 1: - # mb_stop_i = next_frontier.size(0) - # i_nnz_mask = (next_frontier._indices()[0,:] >= mb_start_i) & \ - # (next_frontier._indices()[0,:] < mb_stop_i) - # max_len = max(i_nnz_mask.size(0), max_len) - # max_len_samples = max(mb_stop_i - mb_start_i, max_len_samples) - - # dart_hits_count_recv = torch.cuda.IntTensor(max_len * replication) - # sampled_count_recv = torch.cuda.IntTensor(max_len_samples * replication) - - # if dart_hits_count_select.size(0) < max_len: - # padding = torch.cuda.IntTensor(max_len - dart_hits_count_select.size(0)).fill_(-1) - # print(f"dart.size: {dart_hits_count_select.size()} dartpadding.size: {padding.size()}", flush=True) - # dart_hits_count_select = torch.cat((dart_hits_count_select, padding), dim=0) - # if sampled_count.size(0) < max_len_samples: - # padding = torch.cuda.IntTensor(max_len_samples - sampled_count.size(0)).fill_(-1) - # print(f"sample.size: {sampled_count.size()} samplepadding.size: {padding.size()}", flush=True) - # sampled_count = torch.cat((sampled_count, padding), dim=0) - - dart_lens = [] - sample_lens = [] - for i in range(replication): - mb_start_i = i * int(next_frontier.size(0) / replication) - mb_stop_i = (i + 1) * int(next_frontier.size(0) / replication) - if i == replication - 1: - mb_stop_i = next_frontier.size(0) - i_nnz_mask = (next_frontier._indices()[0,:] >= mb_start_i) & \ - (next_frontier._indices()[0,:] < mb_stop_i) - # max_len = max(i_nnz_mask.size(0), max_len) - dart_lens.append(i_nnz_mask.size(0)) - # max_len_samples = max(mb_stop_i - mb_start_i, max_len_samples) - sample_lens.append(mb_stop_i - mb_start_i) - - dart_hits_count_recv = [] - sample_count_recv = [] - for i in range(replication): - dart_hits_count_recv.append(torch.cuda.IntTensor(dart_lens[i])) - sample_count_recv.append(torch.cuda.IntTensor(sample_lens[i])) - timing_dict["sample-gather-counts-preproc"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - # dist.all_gather_into_tensor(dart_hits_count_recv, dart_hits_count_select, group=row_groups[rank_c]) - # dist.all_gather_into_tensor(sampled_count_recv, sampled_count, group=row_groups[rank_c]) - # dart_hits_count = dart_hits_count_recv[(dart_hits_count_recv != -1)] - # sampled_count = sampled_count_recv[(sampled_count_recv != -1)] - row_procs = list(range(rank_c * replication, (rank_c + 1) * replication)) - dart_hits_count_recv = ring_allgather(dart_hits_count_select, rank, row_procs) - sampled_count_recv = ring_allgather(sampled_count, rank, row_procs) - dart_hits_count = torch.cat(dart_hits_count_recv, dim=0) - sampled_count = torch.cat(sampled_count_recv, dim=0) + 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() From c295033b4c5d9e0d43e7a625bcbf80e6c7e40368 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 7 Oct 2023 16:32:25 -0700 Subject: [PATCH 155/178] Added timing argument to timing functions --- cagnet/samplers/sage.py | 27 +++++++++++++++++------ cagnet/samplers/utils.py | 16 ++++++++++---- examples/gcn_15d.py | 46 ++++++++++++++++++++-------------------- 3 files changed, 55 insertions(+), 34 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index ff5a740..c5678d3 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -16,12 +16,20 @@ timing = True baseline_compare = True -def start_time(timer): - if timing: +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): - if timing: +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) @@ -117,6 +125,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, sa_masks, timing_dict, "sage", timing_arg) + dist.barrier() # adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() @@ -124,6 +133,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, timing_dict, "sage") + dist.barrier() start_time(start_timer) # add explicit 0's to next_frontier to fix the number of rows as bs * sn^l @@ -157,7 +167,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) # current_frontier_select = torch.masked_select(current_frontier.col_indices(), \ # current_frontier.values().bool()).view(current_frontier._nnz(), 1) - timing_dict["frontier-row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) + timing_dict["frontier-row-col-select"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) @@ -186,10 +196,13 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) # adj_matrices[i] = adj_matrix_sample # frontiers[i + 1] = next_frontier_select.clone() - frontiers[i + 1] = next_frontier_select + if i + 1 == n_layers: + frontiers[i + 1] = next_frontier_select + else: + del next_frontier_select current_frontier = next_frontier + timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer)) adj_matrices[i] = adj_matrix_sample.to_sparse_csr() - timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer, barrier=True)) # del next_frontier_select # del current_frontier_select diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index b8bd74a..fcc3ee3 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -16,12 +16,20 @@ timing = True -def start_time(timer): - if timing: +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): - if timing: +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) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 0871195..e10cfb8 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -564,12 +564,12 @@ def main(args, batches=None): dist.send(test_idx_len, dst=dst_gpu) dist.send(train_idx, dst=dst_gpu) dist.send(test_idx, dst=dst_gpu) - print("coalescing", flush=True) - g_loc = g_loc.coalesce().t_() - print("normalizing", flush=True) - g_loc = g_loc.to(device) - features_loc = features_loc.to(device) - edge_count = adj_matrix.size(1) + print("coalescing", flush=True) + g_loc = g_loc.coalesce().t_() + print("normalizing", flush=True) + g_loc = g_loc.to(device) + features_loc = features_loc.to(device) + edge_count = adj_matrix.size(1) else: print(f"loading", flush=True) g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", @@ -849,7 +849,7 @@ def main(args, batches=None): nnz_recv_upperbound = edge_count // (size // args.replication) if epoch >= 1: - start_time(start_timer) + start_time(start_timer, timing_arg=True) print("Sampling", flush=True) torch.cuda.nvtx.range_push("nvtx-sampling") if args.sample_method == "ladies": @@ -870,11 +870,11 @@ def main(args, batches=None): col_groups, args.timing, args.baseline) if epoch >= 1: - model.timings["sample"].append(stop_time(start_timer, stop_timer)) + 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) + start_time(start_timer, timing_arg=True) # g_loc = g_loc.coalesce() print("Training", flush=True) @@ -917,7 +917,7 @@ def main(args, batches=None): adj_row_select_max = min(adj_row_select_max, adj_matrices_bulk[l].size(0)) if epoch >= 1: - start_time(start_inner_timer) + 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] @@ -930,10 +930,10 @@ def main(args, batches=None): 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)) + model.timings["extract-select"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) if epoch >= 1: - start_time(start_inner_timer) + start_time(start_inner_timer, timing_arg=True) if args.sample_method == "ladies": adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ adj_matrix_sample_values, \ @@ -953,18 +953,18 @@ def main(args, batches=None): adj_nnzs.append(adj_sample_values.size(0)) # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") if epoch >= 1: - model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer)) + model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) if epoch >= 1: - start_time(start_inner_timer) + 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() if epoch >= 1: - model.timings["extract-coalesce"].append(stop_time(start_inner_timer, stop_inner_timer)) + 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) + start_time(start_inner_timer, timing_arg=True) torch.cuda.nvtx.range_push("nvtx-selectfeats") if size > 1: src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) @@ -1007,7 +1007,7 @@ def main(args, batches=None): torch.cuda.nvtx.range_pop() # nvtx-selectfeats if epoch >= 1: - model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer)) + model.timings["selectfeats"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) # if epoch >= 1: # start_time(start_inner_timer) @@ -1038,7 +1038,7 @@ def main(args, batches=None): # torch.cuda.nvtx.range_pop() # nvtx-fakemats if epoch >= 1: - start_time(start_inner_timer) + 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)) @@ -1054,7 +1054,7 @@ def main(args, batches=None): 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)) + model.timings["fwd"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) if epoch >= 1: start_time(start_inner_timer) @@ -1068,12 +1068,12 @@ def main(args, batches=None): model.timings["loss"].append(stop_time(start_inner_timer, stop_inner_timer)) if epoch >= 1: - start_time(start_inner_timer) + 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)) + 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) @@ -1081,7 +1081,7 @@ def main(args, batches=None): optimizer.step() torch.cuda.nvtx.range_pop() # nvtx-optstep if epoch >= 1: - model.timings["train"].append(stop_time(start_timer, stop_timer)) + model.timings["train"].append(stop_time(start_timer, stop_timer, timing_arg=True)) if last_batch: # break @@ -1100,7 +1100,7 @@ def main(args, batches=None): # data.val_mask.nonzero().squeeze().size(0), col_groups[0]) acc1 = 0.0 acc3 = 0.0 - if args.timing: + 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"]] From fad31d85db1c5361dc5eb36c949d63d8f4b65adc Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sun, 8 Oct 2023 14:47:48 -0700 Subject: [PATCH 156/178] Replicates graph and does not replicate batches --- cagnet/samplers/sage.py | 11 +- cagnet/samplers/utils.py | 40 ++++- examples/gcn_15d.py | 206 ++++++++++++++++++-------- sparse-extension/sparse_coo_tensor.cu | 53 ++++++- 4 files changed, 228 insertions(+), 82 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index c5678d3..52f9c36 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -76,7 +76,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, gpu = torch.device(f"cuda:{torch.cuda.current_device()}") - batches_expand_rows = torch.arange(mb_count * batch_size, device=gpu) + 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, @@ -85,7 +85,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, batches_expand = sparse_coo_tensor_gpu(batches_expand_idxs, batches._values(), torch.Size([mb_count * batch_size, node_count_total])) - batches_expand = batches_expand.to_sparse_csr() + # batches_expand = batches_expand.to_sparse_csr() # adj_matrix = adj_matrix.to_sparse_csr() current_frontier = batches_expand @@ -116,7 +116,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, torch.Size([current_frontier._nnz(), current_frontier.size(1)])) # nnz = current_frontier._nnz() // mb_count nnz = batch_size * (frontier_size ** i) - current_frontier = current_frontier.to_sparse_csr() + # current_frontier = current_frontier.to_sparse_csr() # Expand batches matrix if baseline_compare: @@ -125,7 +125,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, replication, rank, size, row_groups, col_groups, sa_masks, timing_dict, "sage", timing_arg) + dist.barrier() + # adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: p = p.to_sparse_coo() @@ -170,7 +172,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, timing_dict["frontier-row-col-select"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) + # current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) + 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 diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index fcc3ee3..5da1e69 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1028,11 +1028,35 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # TODO: assume n_layers=1 for now # start_time(start_timer) start_timer.record() - 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_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": + 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] + matc_chunk_row_lens[mata_rows] = adj_row_lens[mata_cols] + matc_chunk_crows = torch.cuda.IntTensor(numerator.size(0) + 1) + matc_chunk_crows[1:] = torch.cumsum(matc_chunk_row_lens, 0, dtype=torch.int32) + matc_chunk_crows[0] = 0 + matc_chunk_cols = torch.cuda.IntTensor(matc_chunk_crows[-1].item()).fill_(0) + rearrange_rows_gpu(numerator._indices()[0,:].int(), + numerator._indices()[1,:].int(), + matc_chunk_crows, + adj_matrix.crow_indices(), + adj_matrix.col_indices(), + matc_chunk_cols) + p_num_rows = torch.repeat_interleave( + torch.arange(matc_chunk_crows.size(0) - 1, device=numerator.device), + matc_chunk_row_lens) + p_num_indices = torch.stack((p_num_rows, matc_chunk_cols)) + p_num_values = torch.cuda.FloatTensor(p_num_indices.size(1)).fill_(1.0) + else: + sa_masks.fill_(False) + # TODO: LADIES code + # 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() @@ -1154,7 +1178,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["sample-filter-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) + # 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() @@ -1274,13 +1298,13 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["select-overflow"] = [] timing_dict["select-iter"] = [] - dist.barrier(group=row_groups[rank_c]) + # 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) # 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]) + # 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() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index e10cfb8..9a9f5a7 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -309,37 +309,37 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep # 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) - # 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) - - 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) - - # input_partitions = torch.split(inputs, math.ceil(float(inputs.size(0)) / (size / replication)), dim=0) + # 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) + # # 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) + + # 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) + + # # 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) @@ -347,12 +347,13 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep input_partitions_fused[-1] = torch.cat(input_partitions[-2:], dim=0) input_partitions = input_partitions_fused - adj_matrix_loc = am_partitions[rank_c] + # 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} 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 + # return inputs_loc, adj_matrix_loc, am_partitions, input_partitions + return inputs_loc, None, None, input_partitions def one5d_partition_mb(rank, size, batches, replication, mb_count): rank_c = rank // replication @@ -467,6 +468,7 @@ def main(args, batches=None): import ogb data = Data() from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time + print(flush=True, end='') 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] @@ -523,31 +525,67 @@ def main(args, batches=None): inputs, num_classes, args.replication, \ args.normalize) + print(f"begin conversion", flush=True) + # adj_matrix = adj_matrix.int() + # adj_matrix_crows = torch.IntTensor(inputs.size(0) + 1).fill_(0) + # ones = torch.IntTensor(adj_matrix.size(1)).fill_(1) + # # adj_matrix_rows = torch.histc(adj_matrix[0,:], bins=inputs.size(0)) + # adj_matrix_rows = torch.IntTensor(inputs.size(0)).fill_(0) + # adj_matrix_rows.scatter_add_(0, adj_matrix[0,:].long(), ones) + # adj_matrix_crows[1:] = torch.cumsum(adj_matrix_rows, dim=0).roll(1) + # adj_matrix_crows[0] = 0 + adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") + adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") + print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) + print(f"adj_matrix_cols: {adj_matrix_cols} dtype: {adj_matrix_cols.dtype} device: {adj_matrix_cols.device}", flush=True) + print(f"stop conversion", flush=True) + # adj_matrix_crows = adj_matrix_crows.int().to(device) + # adj_matrix_cols = adj_matrix[1,:].int().to(device) + adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) + g_loc = torch.sparse_csr_tensor(adj_matrix_crows.int(), adj_matrix_cols.int(), \ + adj_matrix_vals, \ + torch.Size([inputs.size(0), inputs.size(0)])) + g_loc = g_loc.to(device) + print(f"stop copying", flush=True) for i in range(1, args.gpu): print(f"iteration i: {i}", flush=True) 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_() + # g_send = am_partitions[rank_send_row] + g_send = adj_matrix + # 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() + # 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[0] = g_send._nnz() + # g_send_meta[1] = g_send.size(0) + # g_send_meta[2] = g_send.size(1) + 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 + print(f"g_send_meta: {g_send_meta}", flush=True) 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(g_send._indices(), dst=dst_gpu) + # dist.send(g_send._values(), dst=dst_gpu) + # print(f"i: {i} g_send: {g_send} dtype: {g_send.dtype}", flush=True) + # dist.send(g_send, dst=dst_gpu) + print(f"i: {i} g_loc.crows: {g_loc.crow_indices()} dtype: {g_loc.crow_indices().dtype}", flush=True) + dist.send(g_loc.crow_indices(), dst=dst_gpu) + print(f"i: {i} g_loc.cols: {g_loc.col_indices()} dtype: {g_loc.col_indices().dtype}", flush=True) + dist.send(g_loc.col_indices(), dst=dst_gpu) + print(f"i: {i} features_send.size: {features_send.size()} dtype: {features_send.dtype}", flush=True) dist.send(features_send, dst=dst_gpu) + print(f"send data.y.size: {data.y.size()} inputs.size: {inputs.size()}", flush=True) dist.send(data.y.float(), dst=dst_gpu) # torch.save(g_send_meta, f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_meta.pt") # torch.save(g_send._indices(), f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_indices.pt") @@ -565,11 +603,11 @@ def main(args, batches=None): dist.send(train_idx, dst=dst_gpu) dist.send(test_idx, dst=dst_gpu) print("coalescing", flush=True) - g_loc = g_loc.coalesce().t_() + # g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) g_loc = g_loc.to(device) features_loc = features_loc.to(device) - edge_count = adj_matrix.size(1) + edge_count = g_loc.col_indices().size(0) else: print(f"loading", flush=True) g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", @@ -593,7 +631,7 @@ def main(args, batches=None): g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, torch.Size([g_loc_meta[1], g_loc_meta[2]])) print("normalizing", flush=True) - edge_count = adj_matrix.size(1) + edge_count = g_loc.col_indices().size(0) # features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ # args.replication, args.normalize) @@ -617,28 +655,55 @@ def main(args, batches=None): 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_indices = torch.cuda.IntTensor(2, g_loc_meta[0].item()) + # dist.recv(g_loc_indices, src=src_gpu) + # torch.cuda.synchronize() + # print(f"recv indices", flush=True) + + g_loc_crows = torch.cuda.IntTensor(g_loc_meta[1].item() + 1) + dist.recv(g_loc_crows, src=src_gpu) + torch.cuda.synchronize() + print(f"recv crows", flush=True) - g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) - dist.recv(g_loc_values, src=src_gpu) + g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) + dist.recv(g_loc_cols, src=src_gpu) + torch.cuda.synchronize() + print(f"recv cols", flush=True) + + # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + # dist.recv(g_loc_values, src=src_gpu) num_features = g_loc_meta[4].item() num_classes = g_loc_meta[5].item() + rows_per_proc = g_loc_meta[1].item() // (size // args.replication) if rank_row < proc_row - 1: - inputs = torch.cuda.FloatTensor(g_loc_meta[1].item(), num_features) + inputs = torch.cuda.FloatTensor(rows_per_proc, num_features) else: - rows = g_loc_meta[1].item() + (g_loc_meta[2].item() - proc_row * g_loc_meta[1].item()) + rows = g_loc_meta[1].item() - rows_per_proc * (proc_row - 1) inputs = torch.cuda.FloatTensor(rows, num_features) dist.recv(inputs, src=src_gpu) + torch.cuda.synchronize() + print(f"recv inputs", flush=True) # data = Data() data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + print(f"data.y.size: {data.y.size()}", flush=True) dist.recv(data.y, src=src_gpu) + torch.cuda.synchronize() + print(f"recv datay", flush=True) data.y = data.y.long() - - g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + print(f"cast datay", flush=True) + + # g_loc_values = torch.cuda.FloatTensor(g_loc_indices.size(1)).fill_(1.0) + # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + # torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc_values = torch.cuda.FloatTensor(g_loc_cols.size(0)).fill_(1.0) + g_loc_crows = g_loc_crows.cpu() + g_loc_cols = g_loc_cols.cpu() + g_loc_values = g_loc_values.cpu() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc = g_loc.to(device) edge_count = g_loc_meta[3].item() if args.dataset == "ogbn-papers100M": @@ -687,10 +752,19 @@ def main(args, batches=None): features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ args.replication, args.normalize) print("coalescing", flush=True) - g_loc = g_loc.coalesce().t_() + # g_loc = g_loc.coalesce().t_() + 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,:], 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)])) print("normalizing", flush=True) g_loc = g_loc.to(device) - g_loc = g_loc.double() + # g_loc = g_loc.double() # features_loc = inputs features_loc = features_loc.to(device) # del inputs @@ -797,7 +871,7 @@ def main(args, batches=None): # if batches is None: torch.manual_seed(epoch) vertex_perm = torch.randperm(train_nid.size(0)) - batches_all = train_nid[vertex_perm] + batches_all = train_nid[vertex_perm].int() torch.cuda.nvtx.range_pop() # construct-batches # torch.cuda.nvtx.range_push("nvtx-construct-batches") @@ -821,13 +895,13 @@ def main(args, batches=None): args.batch_size = batch_count - b torch.cuda.nvtx.range_push("nvtx-partition-batches") batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) - batches_loc = one5d_partition_mb(rank, size, batches, args.replication, args.n_bulkmb) + batches_loc = one5d_partition_mb(rank, size, batches, 1, args.n_bulkmb) - batches_indices_rows = torch.arange(batches_loc.size(0)).to(device) + 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.DoubleTensor(batches_loc.size(1) * batches_loc.size(0), device=device).fill_(1.0) + 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) @@ -843,8 +917,8 @@ def main(args, batches=None): # for sa-spgemm # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) - nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._nnz()) - nnz_row_masks.fill_(0) + # 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) @@ -862,10 +936,12 @@ def main(args, batches=None): row_groups, col_groups, args.timing) elif args.sample_method == "sage": # nnz_row_masks.fill_(False) + nnz_row_masks = None 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, \ - args.replication, nnz_row_masks, + 1, nnz_row_masks, + # args.replication, nnz_row_masks, rank, size, row_groups, col_groups, args.timing, args.baseline) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index d580505..e2a740a 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2586,7 +2586,7 @@ void sort_dst_proc_gpu(const at::Tensor& vtxs, const at::Tensor& src_vtx_sort, c CHECK_ERROR("sort_dst_proc error") } -__global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, long *matb_crows, +__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; @@ -2605,9 +2605,9 @@ __global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows } } -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) { +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; @@ -2618,7 +2618,7 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols return; } - RearrangeRows<<>>(mata_rows.data(), + RearrangelRows<<>>(mata_rows.data(), mata_cols.data(), matc_crows.data(), matb_crows.data(), @@ -2629,6 +2629,49 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols CHECK_ERROR("rearrange_rows error") } +__global__ void RearrangeRows(int *mata_rows, int *mata_cols, int *matc_crows, int *matb_crows, + int *matb_cols, int *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) { + int row = mata_rows[i]; + int col = mata_cols[i]; + int dst_idx = matc_crows[row]; + int src_idx = matb_crows[col]; + int 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 rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols, + const at::Tensor& matc_crows, const at::Tensor& matb_crows, + const at::Tensor& matb_cols, const at::Tensor& matc_cols) { + + + int BLOCK_SIZE = 256; + int BLOCK_COUNT = std::ceil(mata_rows.sizes()[0] / ((float) BLOCK_SIZE)); + BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); + + if (mata_rows.sizes()[0] == 0) { + return; + } + + RearrangeRows<<>>(mata_rows.data(), + mata_cols.data(), + matc_crows.data(), + matb_crows.data(), + matb_cols.data(), + matc_cols.data(), + mata_rows.sizes()[0]); + + CHECK_ERROR("rearrange_rows error") +} + // __global__ void ReduceSum(long *matc_crows, long *mata_crows, long *matb_crows, long *matc_cols, // long *mata_cols, long *matb_cols, int row_count) { // From dd1efad7e53f8bda9f9ed790d0eab32711aed862 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sun, 8 Oct 2023 17:21:34 -0700 Subject: [PATCH 157/178] Removed unnecessary barriers --- cagnet/samplers/sage.py | 4 ++-- examples/gcn_15d.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 52f9c36..b2dbf3c 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -126,7 +126,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, sa_masks, timing_dict, "sage", timing_arg) - dist.barrier() + # dist.barrier() # adj_matrix = adj_matrix.to_sparse_coo() if p.layout == torch.sparse_csr: @@ -135,7 +135,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, next_frontier = sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, rank, size, row_groups, col_groups, timing_dict, "sage") - dist.barrier() + # dist.barrier() start_time(start_timer) # add explicit 0's to next_frontier to fix the number of rows as bs * sn^l diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 9a9f5a7..80a06b4 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -467,8 +467,9 @@ def main(args, batches=None): elif args.dataset.startswith("ogb"): import ogb data = Data() + print(f"before loading ogb", flush=True) from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time - print(flush=True, end='') + print(f"after loading ogb", flush=True) 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] From 07f681ccd2e64affdb277cdf38d066ea907cbb58 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 9 Oct 2023 18:58:15 -0700 Subject: [PATCH 158/178] Converted doubles to floats to save sapce --- cagnet/samplers/sage.py | 13 ++-- cagnet/samplers/utils.py | 48 +++++++------- examples/gcn_15d.py | 42 ++++++++----- sparse-extension/sparse_coo_tensor.cu | 90 ++++++++++++++++++++------- 4 files changed, 125 insertions(+), 68 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index b2dbf3c..2f9d717 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -2,6 +2,7 @@ import torch import torch.distributed as dist import torch_sparse +import numpy as np from collections import defaultdict from cagnet.samplers.utils import * @@ -39,7 +40,7 @@ def stop_time(start_timer, stop_timer, barrier=False, timing_arg=None): else: return 0.0 -def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, n_layers, n_darts, \ +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): @@ -64,11 +65,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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 mb mtx for layer j @@ -115,7 +111,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, 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 * (frontier_size ** i) + nnz = batch_size * int(np.prod(frontier_sizes[:i], dtype=int)) # current_frontier = current_frontier.to_sparse_csr() # Expand batches matrix @@ -132,6 +129,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_total, if p.layout == torch.sparse_csr: p = p.to_sparse_coo() + 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") diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 5da1e69..0200cdc 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1038,19 +1038,21 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication mata_cols = numerator._indices()[1,:] adj_row_lens = adj_matrix.crow_indices()[1:] - adj_matrix.crow_indices()[:-1] matc_chunk_row_lens[mata_rows] = adj_row_lens[mata_cols] - matc_chunk_crows = torch.cuda.IntTensor(numerator.size(0) + 1) - matc_chunk_crows[1:] = torch.cumsum(matc_chunk_row_lens, 0, dtype=torch.int32) + 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,:].int(), - numerator._indices()[1,:].int(), + rearrange_rows_gpu(numerator._indices()[0,:], + numerator._indices()[1,:], matc_chunk_crows, - adj_matrix.crow_indices(), - adj_matrix.col_indices(), + adj_matrix.crow_indices().long(), + adj_matrix.col_indices().long(), matc_chunk_cols) p_num_rows = torch.repeat_interleave( - torch.arange(matc_chunk_crows.size(0) - 1, device=numerator.device), - matc_chunk_row_lens) + 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_values = torch.cuda.FloatTensor(p_num_indices.size(1)).fill_(1.0) else: @@ -1065,24 +1067,25 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication timing_dict["probability-spgemm"].append(time_taken) start_time(start_timer) - p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) 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() 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_num_values = p_num_values.double() + p_den = torch.cuda.FloatTensor(numerator.size(0)).fill_(0) # scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) # frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) - p_den.scatter_add_(0, p_num_indices[0, :], p_num_values) + 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()) - normalize_gpu(p_num_values, p_den, p_num_indices[0, :], p_num_values.size(0)) + 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, p_num_values, torch.Size([numerator.size(0), node_count_total])) + p = sparse_coo_tensor_gpu(p_num_indices.long(), p_num_values, torch.Size([numerator.size(0), node_count_total])) timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) return p @@ -1104,10 +1107,10 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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 + # 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 next_frontier = torch.sparse_coo_tensor(indices=p._indices(), values=torch.cuda.LongTensor(p._nnz()).fill_(0), @@ -1135,7 +1138,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True - p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(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) @@ -1160,7 +1164,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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.DoubleTensor(n_darts_col * p.size(0)).uniform_() + dart_values = torch.cuda.FloatTensor(n_darts_col * p.size(0)).uniform_() timing_dict["sample-gen-darts"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) @@ -1207,6 +1212,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(select_start_timer) + print(f"overflow: {overflow}", flush=True) if rank_col == 0 and overflowed_minibatches: while overflowed_minibatches: start_time(select_iter_start_timer) @@ -1338,8 +1344,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) - # print(f"iter_count: {iter_count}") - # print(f"selection_iter_count: {selection_iter_count}") + print(f"iter_count: {iter_count}") + print(f"selection_iter_count: {selection_iter_count}") return next_frontier diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 80a06b4..89683ba 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -85,6 +85,7 @@ def forward(self, graphs, inputs, epoch): for l, layer in enumerate(self.layers): graphs[l].t_() nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv + # print(f"nnz_index.size: {nnz_index.size()}", flush=True) edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv # edge_index = graphs[l] # SAGEConvfake # edge_index, _, size = graphs[l] # quiver @@ -827,7 +828,8 @@ def main(args, batches=None): model = model.to(device) # use optimizer - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) + # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) dur = [] # torch.manual_seed(0) @@ -886,6 +888,7 @@ def main(args, batches=None): 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 @@ -914,7 +917,12 @@ def main(args, batches=None): 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 * 2 + 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) + print(f"n_darts: {args.n_darts}", flush=True) # for sa-spgemm # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) @@ -975,22 +983,20 @@ def main(args, batches=None): batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) # src_vtxs = frontiers[-1][i].view(-1) - src_select_size = args.batch_size * (args.samp_num ** (args.n_layers - 1)) - src_select_min = i * args.batch_size * (args.samp_num ** (args.n_layers - 1)) - src_select_max = (i + 1) * args.batch_size * (args.samp_num ** (args.n_layers - 1)) + # src_select_min = i * args.batch_size * (args.samp_num ** (args.n_layers - 1)) + # src_select_max = (i + 1) * args.batch_size * (args.samp_num ** (args.n_layers - 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 += src_select_size * rank_col * rank_n_bulkmb - # src_select_max += src_select_size * rank_col * rank_n_bulkmb src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) # print(f"src_vtxs: {src_vtxs} src_vtxs.size: {src_vtxs.size()} zero_count: {(src_vtxs == 0).sum()}", flush=True) adjs = [None] * args.n_layers for l in range(args.n_layers): - adj_select_size = args.batch_size * (args.samp_num ** l) - 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 += adj_select_size * rank_col * rank_n_bulkmb - # adj_row_select_max += adj_select_size * rank_col * rank_n_bulkmb + # adj_row_select_min = i * args.batch_size * (args.samp_num ** l) + # adj_row_select_max = (i + 1) * args.batch_size * (args.samp_num ** l) + adj_row_select_min = i * args.batch_size * np.prod(args.samp_num[:l], dtype=int) + adj_row_select_max = (i + 1) * args.batch_size * np.prod(args.samp_num[:l], dtype=int) adj_row_select_max = min(adj_row_select_max, adj_matrices_bulk[l].size(0)) if epoch >= 1: @@ -1025,8 +1031,10 @@ def main(args, batches=None): 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 * (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)) # 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: @@ -1154,6 +1162,7 @@ def main(args, batches=None): for W in model.parameters(): dist.all_reduce(W.grad) + W.grad /= size torch.cuda.nvtx.range_push("nvtx-optstep") optimizer.step() torch.cuda.nvtx.range_pop() # nvtx-optstep @@ -1252,8 +1261,8 @@ def main(args, batches=None): 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("--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", @@ -1287,6 +1296,7 @@ def main(args, batches=None): parser.add_argument('--baseline', action="store_true", help='whether to avoid col selection for baseline comparison') args = parser.parse_args() + args.samp_num = [int(i) for i in args.samp_num.split('-')] print(args) main(args) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index e2a740a..98e6c48 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1330,8 +1330,24 @@ using namespace at::sparse; // } // cudaThreadSynchronize(); // } +__device__ int binary_searchd(double *arr, double val, int imin, int imax) { + + int ans = 0; + while (imax >= imin) { + int imid = (imin + imax) / 2; -__device__ int binary_searchf(double *arr, double val, int imin, int imax) { + 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) { @@ -2085,7 +2101,8 @@ void compute_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart computation error") } -__global__ void ComputeDarts1D(double *dart_values, double *p_rowsum, double *ps_p_rowsum, +// __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, int n_darts, int mb_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; @@ -2107,9 +2124,12 @@ void compute_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& p_rows int BLOCK_COUNT = std::ceil((n_darts * mb_count) / ((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(), + ComputeDarts1D<<>>(dart_values.data(), + p_rowsum.data(), + ps_p_rowsum.data(), n_darts, mb_count); CHECK_ERROR("dart1d computation error") @@ -2189,7 +2209,7 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_values, +__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, int *h_values, int dart_count, int nnz) { int id = blockIdx.x * blockDim.x + threadIdx.x; @@ -2215,8 +2235,10 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, 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(), + ThrowDarts1D<<>>(dart_values.data(), + ps_p_values.data(), h_values.data(), dart_count, nnz); @@ -2230,7 +2252,8 @@ __global__ void ThrowDartsSelect(double *dart_select, double *ps_dart_hits_inv, 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_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); } } @@ -2258,7 +2281,26 @@ void throw_darts_select_gpu(const at::Tensor& dart_select, CHECK_ERROR("selection dart throwing error") } -__global__ void Normalize(double *output, double *input, long *index, int len) { +__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(float *output, float *input, long *index, int len) { long id = blockIdx.x * blockDim.x + threadIdx.x; long stride = blockDim.x * gridDim.x; @@ -2274,7 +2316,7 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - Normalize<<>>(output.data(), input.data(), index.data(), len); + Normalize<<>>(output.data(), input.data(), index.data(), len); } __global__ void NormalizeBatch(float *output, int *input, int output_rows, int output_cols) { @@ -2629,21 +2671,21 @@ void rearrangel_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_col CHECK_ERROR("rearrange_rows error") } -__global__ void RearrangeRows(int *mata_rows, int *mata_cols, int *matc_crows, int *matb_crows, - int *matb_cols, int *matc_cols, int nnz_count) { +__global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, long *matb_crows, + long *matb_cols, int *matc_cols, long nnz_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; for (int i = id; i < nnz_count; i += stride) { - int row = mata_rows[i]; - int col = mata_cols[i]; - int dst_idx = matc_crows[row]; - int src_idx = matb_crows[col]; - int row_len = matb_crows[col + 1] - matb_crows[col]; + 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]; + matc_cols[dst_idx + j] = (int)(matb_cols[src_idx + j]); } } } @@ -2661,11 +2703,11 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols return; } - RearrangeRows<<>>(mata_rows.data(), - mata_cols.data(), - matc_crows.data(), - matb_crows.data(), - matb_cols.data(), + RearrangeRows<<>>(mata_rows.data(), + mata_cols.data(), + matc_crows.data(), + matb_crows.data(), + matb_cols.data(), matc_cols.data(), mata_rows.sizes()[0]); From 60b3d4d6b2f974c4984ffed76c0ffb4b09ce20b6 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 9 Oct 2023 21:20:28 -0700 Subject: [PATCH 159/178] Cached feature vtx 0 and removed nonzero() call in forward() --- cagnet/samplers/utils.py | 1 - examples/gcn_15d.py | 33 +++++++++++++++++++++++++++------ 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 0200cdc..47730c0 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1212,7 +1212,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(select_start_timer) - print(f"overflow: {overflow}", flush=True) if rank_col == 0 and overflowed_minibatches: while overflowed_minibatches: start_time(select_iter_start_timer) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 89683ba..a1940cf 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -84,9 +84,9 @@ def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): graphs[l].t_() - nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv - # print(f"nnz_index.size: {nnz_index.size()}", flush=True) - edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv + # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv + edge_index = graphs[l]._indices() # edge_index = graphs[l] # SAGEConvfake # edge_index, _, size = graphs[l] # quiver h = self.layers[l](h, edge_index) # SAGEConv @@ -548,6 +548,7 @@ def main(args, batches=None): adj_matrix_vals, \ torch.Size([inputs.size(0), inputs.size(0)])) g_loc = g_loc.to(device) + 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) @@ -587,6 +588,8 @@ def main(args, batches=None): dist.send(g_loc.col_indices(), dst=dst_gpu) print(f"i: {i} features_send.size: {features_send.size()} dtype: {features_send.dtype}", flush=True) dist.send(features_send, dst=dst_gpu) + print(f"i: {i} features_zero_row.size: {features_zero_row.size()} dtype: {features_zero_row.dtype}", flush=True) + dist.send(features_zero_row, dst=dst_gpu) print(f"send data.y.size: {data.y.size()} inputs.size: {inputs.size()}", flush=True) dist.send(data.y.float(), dst=dst_gpu) # torch.save(g_send_meta, f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_meta.pt") @@ -687,6 +690,11 @@ def main(args, batches=None): torch.cuda.synchronize() print(f"recv inputs", flush=True) + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + # data = Data() data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) print(f"data.y.size: {data.y.size()}", flush=True) @@ -768,6 +776,7 @@ def main(args, batches=None): 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 @@ -1052,8 +1061,12 @@ def main(args, batches=None): start_time(start_inner_timer, timing_arg=True) torch.cuda.nvtx.range_push("nvtx-selectfeats") if size > 1: + src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) og_idxs = torch.cuda.LongTensor(src_vtxs.size(0)) + # src_vtxs_nnz = src_vtxs[src_vtxs_nnz_mask] + # src_vtxs_sort = torch.cuda.LongTensor(src_vtxs_nnz.size(0)) + # og_idxs = torch.cuda.LongTensor(src_vtxs_nnz.size(0)) tally.fill_(0) # node_per_row = int(math.ceil(node_count / (size / args.replication))) @@ -1063,6 +1076,11 @@ def main(args, batches=None): # 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 = [] @@ -1080,13 +1098,16 @@ def main(args, batches=None): # 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.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[og_idxs] = output_features + #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 else: features_batch = features_loc[src_vtxs] From 914c196ed479bbddf48652f0b1c7018a29c8df79 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 11 Oct 2023 20:05:35 -0700 Subject: [PATCH 160/178] Dart count is computed per iteration --- cagnet/samplers/utils.py | 34 +++++++++-- examples/gcn_15d.py | 87 ++++++++++++++++----------- sparse-extension/sparse_coo_tensor.cu | 36 ++++++++--- 3 files changed, 108 insertions(+), 49 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 47730c0..0ab77f7 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1110,7 +1110,11 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, 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 + + # 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 next_frontier = torch.sparse_coo_tensor(indices=p._indices(), values=torch.cuda.LongTensor(p._nnz()).fill_(0), @@ -1128,7 +1132,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, zero_idxs = p._indices()[0, (p._values() == 0)].unique() frontier_nnz_sizes[zero_idxs] = 0 - + timing_dict["sample-pre-loop"].append(stop_time(start_timer, stop_timer)) sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) @@ -1138,6 +1142,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True + print(f"frontier_size: {frontier_size}", flush=True) # p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) p_rowsum = torch.cuda.FloatTensor(p.size(0)).fill_(0) while underfull_minibatches: @@ -1162,31 +1167,43 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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 + dart_count_row = frontier_size - sampled_count + n_darts_col = dart_count_row.sum().item() + 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 * p.size(0)).uniform_() + dart_values = torch.cuda.FloatTensor(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)) - + print(f"dart_count_row: {dart_count_row} max: {dart_count_row.max()}", flush=True) 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, 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.size(0), p._nnz()) + n_darts_col, p._nnz()) + # n_darts_col * p.size(0), p._nnz()) # n_darts * mb_count, p._nnz()) timing_dict["sample-filter-darts"].append(stop_time(start_timer, stop_timer)) + print(f"dart_hits_count.sum: {dart_hits_count.sum()} n_darts_col: {n_darts_col} size: {sampled_count.size()}", flush=True) start_time(start_timer) # dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) + print(f"before dart_hits_count.nnz: {dart_hits_count.nonzero().squeeze().size()} next_frontier.vals.nnz: {next_frontier._values().nonzero().squeeze().size()}", flush=True) next_frontier_values = torch.logical_or( dart_hits_count, next_frontier._values().int()).int() + print(f"after next_frontier_values.nnz: {next_frontier_values.nonzero().squeeze().size()}", flush=True) # next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), # values=next_frontier_values, # size=(p.size(0), node_count_total)) @@ -1196,6 +1213,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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) next_frontier_nnzmask = next_frontier_values.nonzero().squeeze() @@ -1205,6 +1223,9 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, sampled_count.scatter_add_(0, next_frontier_nnzidxs, next_frontier_nnzvals) timing_dict["sample-count-samples"].append(stop_time(start_timer, stop_timer)) + print(f"diff: {sampled_count - sampled_count_old} diff_max: {(sampled_count - sampled_count_old).max()} frontier_size: {frontier_size} sum: {(sampled_count - sampled_count_old).sum()}", flush=True) + del sampled_count_old + start_time(start_timer) overflow = torch.clamp(sampled_count - frontier_size, min=0).int() overflowed_minibatches = (overflow > 0).any() @@ -1212,6 +1233,7 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(select_start_timer) + print(f"total_overflow: {overflow.sum()} sampled_count.max: {sampled_count.max()} sum: {sampled_count.sum()}", flush=True) if rank_col == 0 and overflowed_minibatches: while overflowed_minibatches: start_time(select_iter_start_timer) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index a1940cf..46544fa 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -466,33 +466,43 @@ def main(args, batches=None): data = data.to(device) data.y = data.y.to(device) elif args.dataset.startswith("ogb"): - import ogb + # import ogb data = Data() - print(f"before loading ogb", flush=True) - from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time - print(f"after loading ogb", flush=True) + # 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 + # 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) + 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) - # 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) + print(f"data.y: {data.y} data.y.dtype {data.y.dtype}", flush=True) + train_idx = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/index/train_idx.pt") + train_idx = train_idx.to(device) + print(f"train_idx: {train_idx} train_idx.dtype {train_idx.dtype}", flush=True) + inputs = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") + print(f"inputs: {inputs} inputs.dtype {inputs.dtype}", flush=True) + num_features = inputs.size(1) + num_classes = 172 # print(f"adj_matrix.size(): {adj_matrix.size()}") # print(f"adj_matrix: {adj_matrix}") @@ -507,7 +517,11 @@ def main(args, batches=None): partitioning = Partitioning.ONE5D + print(f"before sync", flush=True) + torch.cuda.synchronize() + print(f"before get_proc_group", flush=True) row_groups, col_groups = get_proc_groups(rank, size, args.replication) + print(f"after get_proc_group", flush=True) proc_row = size // args.replication rank_row = rank // args.replication @@ -520,7 +534,7 @@ def main(args, batches=None): if rank % args.gpu == 0: if True or args.dataset != "ogbn-papers100M": inputs = inputs.to(torch.device("cpu")) - adj_matrix = adj_matrix.to(torch.device("cpu")) + # adj_matrix = adj_matrix.to(torch.device("cpu")) # 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, @@ -563,7 +577,8 @@ def main(args, batches=None): # g_send = g_send.double() features_send = input_partitions[rank_send_row].to(device) - edge_count = adj_matrix.size(1) + # edge_count = adj_matrix.size(1) + edge_count = g_loc.col_indices().size(0) g_send_meta = torch.cuda.LongTensor(6) @@ -602,11 +617,11 @@ def main(args, batches=None): 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)) + # 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(test_idx_len, dst=dst_gpu) dist.send(train_idx, dst=dst_gpu) - dist.send(test_idx, dst=dst_gpu) + # dist.send(test_idx, dst=dst_gpu) print("coalescing", flush=True) # g_loc = g_loc.coalesce().t_() print("normalizing", flush=True) @@ -718,14 +733,14 @@ def main(args, batches=None): if args.dataset == "ogbn-papers100M": train_idx_len = torch.cuda.LongTensor(1).fill_(0) - test_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) + # 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) + # 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) + # dist.recv(test_idx, src=src_gpu) else: print(f"loading", flush=True) # Load data to CPU instead of GPU bc multiple processes spawn for some reason, @@ -785,7 +800,7 @@ def main(args, batches=None): print(f"features_loc.size: {features_loc.size()}", flush=True) dist.barrier() - if rank == 0: + if rank == 0 and args.dataset != "ogbn-papers100M": adj_matrix = adj_matrix.cpu() inputs = inputs.cpu() @@ -806,9 +821,9 @@ def main(args, batches=None): print(f"test_nid.size: {test_nid.size()}") else: train_nid = train_idx - test_nid = test_idx + # test_nid = test_idx print(f"train_nid.size: {train_nid.size()}") - print(f"test_nid.size: {test_nid.size()}") + # print(f"test_nid.size: {test_nid.size()}") # # # adj_matrix = None # Uncomment bottom for testing # adj_matrix = adj_matrix.cuda() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 98e6c48..658477e 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1399,6 +1399,23 @@ __device__ long binary_searchl(long *arr, long val, long imin, long imax) { 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) { @@ -2103,25 +2120,29 @@ void compute_darts_gpu(const at::Tensor& dartx_values, // __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, - int n_darts, int mb_count) { + 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 < dart_count; i += stride) { - int row = i / n_darts; - // dart_values[i] += (double)row; + // 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); + if (row >= mb_count) { + printf("error row %ld mb_count %d\n", row, 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, int n_darts, int mb_count) { + 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 * mb_count) / ((float) BLOCK_SIZE)); + int BLOCK_COUNT = std::ceil(n_darts / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); // ComputeDarts1D<<>>(dart_values.data(), @@ -2130,6 +2151,7 @@ void compute_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& p_rows 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") From 4e583da60e9e78abd371211996c68c0a094332ac Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 11 Oct 2023 20:50:52 -0700 Subject: [PATCH 161/178] Fixed precision errors with some doubles. Papers sampling time 2x'd --- cagnet/samplers/utils.py | 13 ++++++----- sparse-extension/sparse_coo_tensor.cu | 32 +++++++++++++-------------- 2 files changed, 23 insertions(+), 22 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 0ab77f7..a5241ef 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1054,7 +1054,8 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication device=numerator.device), matc_chunk_row_lens) p_num_indices = torch.stack((p_num_rows, matc_chunk_cols)) - p_num_values = torch.cuda.FloatTensor(p_num_indices.size(1)).fill_(1.0) + # 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) else: sa_masks.fill_(False) # TODO: LADIES code @@ -1074,7 +1075,8 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication 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.FloatTensor(numerator.size(0)).fill_(0) + p_den = torch.cuda.DoubleTensor(numerator.size(0)).fill_(0) # scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) # frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) p_den.scatter_add_(0, p_num_indices[0, :].long(), p_num_values) @@ -1143,8 +1145,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, underfull_minibatches = True print(f"frontier_size: {frontier_size}", flush=True) - # p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) - p_rowsum = torch.cuda.FloatTensor(p.size(0)).fill_(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) @@ -1175,7 +1177,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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.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)) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 658477e..f62cdeb 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2118,8 +2118,8 @@ void compute_darts_gpu(const at::Tensor& dartx_values, 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, +__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; @@ -2129,9 +2129,6 @@ __global__ void ComputeDarts1D(float *dart_values, float *p_rowsum, float *ps_p_ 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); - if (row >= mb_count) { - printf("error row %ld mb_count %d\n", row, mb_count); - } dart_values[i] *= p_rowsum[row]; dart_values[i] += ps_p_rowsum[row]; } @@ -2145,12 +2142,12 @@ void compute_darts1d_gpu(const at::Tensor& dart_values, const at::Tensor& p_rows 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(), + 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); @@ -2231,14 +2228,15 @@ void throw_darts_gpu(const at::Tensor& dartx_values, CHECK_ERROR("dart throwing error") } -__global__ void ThrowDarts1D(float *dart_values, float *ps_p_values, int *h_values, +__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_searchf(ps_p_values, dart_values[i], 0, nnz - 1); + int vtx = binary_searchd(ps_p_values, dart_values[i], 0, nnz - 1); // if (vtx < 0 || vtx >= nnz) { // printf("error i: %d vtx: %d nnz: %d\n", i, vtx, nnz); // } @@ -2259,8 +2257,8 @@ void throw_darts1d_gpu(const at::Tensor& dart_values, // ThrowDarts1D<<>>(dart_values.data(), // ps_p_values.data(), - ThrowDarts1D<<>>(dart_values.data(), - ps_p_values.data(), + ThrowDarts1D<<>>(dart_values.data(), + ps_p_values.data(), h_values.data(), dart_count, nnz); @@ -2322,7 +2320,7 @@ void normalized_gpu(const at::Tensor& output, const at::Tensor& input, const at: Normalized<<>>(output.data(), input.data(), index.data(), len); } -__global__ void Normalize(float *output, float *input, long *index, int 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; @@ -2338,7 +2336,7 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: int BLOCK_COUNT = std::ceil(len / ((float) BLOCK_SIZE)); BLOCK_COUNT = std::min(BLOCK_COUNT, 65535); - Normalize<<>>(output.data(), input.data(), index.data(), len); + Normalize<<>>(output.data(), input.data(), index.data(), len); } __global__ void NormalizeBatch(float *output, int *input, int output_rows, int output_cols) { From 58994a119d96b0bb339a6651a72824b7bf121ac6 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 11 Oct 2023 23:02:39 -0700 Subject: [PATCH 162/178] Cleanup --- cagnet/samplers/utils.py | 10 ---------- examples/gcn_15d.py | 3 --- 2 files changed, 13 deletions(-) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index a5241ef..e6f1eb9 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1144,7 +1144,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # underfull_minibatches = (sampled_count < frontier_size).any() underfull_minibatches = True - print(f"frontier_size: {frontier_size}", flush=True) p_rowsum = torch.cuda.DoubleTensor(p.size(0)).fill_(0) # p_rowsum = torch.cuda.FloatTensor(p.size(0)).fill_(0) while underfull_minibatches: @@ -1182,7 +1181,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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)) - print(f"dart_count_row: {dart_count_row} max: {dart_count_row.max()}", flush=True) start_time(start_timer) # compute_darts1d_gpu(dart_values, n_darts, mb_count) # compute_darts1d_gpu(dart_values, n_darts_col, mb_count) @@ -1198,15 +1196,12 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # n_darts_col * p.size(0), p._nnz()) # n_darts * mb_count, p._nnz()) timing_dict["sample-filter-darts"].append(stop_time(start_timer, stop_timer)) - print(f"dart_hits_count.sum: {dart_hits_count.sum()} n_darts_col: {n_darts_col} size: {sampled_count.size()}", flush=True) start_time(start_timer) # dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) - print(f"before dart_hits_count.nnz: {dart_hits_count.nonzero().squeeze().size()} next_frontier.vals.nnz: {next_frontier._values().nonzero().squeeze().size()}", flush=True) next_frontier_values = torch.logical_or( dart_hits_count, next_frontier._values().int()).int() - print(f"after next_frontier_values.nnz: {next_frontier_values.nonzero().squeeze().size()}", flush=True) # next_frontier_tmp = torch.sparse_coo_tensor(indices=next_frontier._indices(), # values=next_frontier_values, # size=(p.size(0), node_count_total)) @@ -1226,7 +1221,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, sampled_count.scatter_add_(0, next_frontier_nnzidxs, next_frontier_nnzvals) timing_dict["sample-count-samples"].append(stop_time(start_timer, stop_timer)) - print(f"diff: {sampled_count - sampled_count_old} diff_max: {(sampled_count - sampled_count_old).max()} frontier_size: {frontier_size} sum: {(sampled_count - sampled_count_old).sum()}", flush=True) del sampled_count_old start_time(start_timer) @@ -1236,7 +1230,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, start_time(select_start_timer) - print(f"total_overflow: {overflow.sum()} sampled_count.max: {sampled_count.max()} sum: {sampled_count.sum()}", flush=True) if rank_col == 0 and overflowed_minibatches: while overflowed_minibatches: start_time(select_iter_start_timer) @@ -1368,9 +1361,6 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) - print(f"iter_count: {iter_count}") - print(f"selection_iter_count: {selection_iter_count}") - return next_frontier def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 46544fa..2949336 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -517,11 +517,8 @@ def main(args, batches=None): partitioning = Partitioning.ONE5D - print(f"before sync", flush=True) torch.cuda.synchronize() - print(f"before get_proc_group", flush=True) row_groups, col_groups = get_proc_groups(rank, size, args.replication) - print(f"after get_proc_group", flush=True) proc_row = size // args.replication rank_row = rank // args.replication From 82bda7f28509ccd7c9658f7196b83f0f150972d3 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 12 Oct 2023 19:37:40 -0700 Subject: [PATCH 163/178] Commented code for training time opt --- examples/gcn_15d.py | 66 ++++++++++++++++++++++++++++++++++++++------- 1 file changed, 57 insertions(+), 9 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 2949336..17ada3c 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -84,8 +84,6 @@ def forward(self, graphs, inputs, epoch): h = inputs for l, layer in enumerate(self.layers): graphs[l].t_() - # nnz_index = graphs[l]._values().nonzero().squeeze() # SAGEConv - # edge_index = graphs[l]._indices()[:, nnz_index] # SAGEConv edge_index = graphs[l]._indices() # edge_index = graphs[l] # SAGEConvfake # edge_index, _, size = graphs[l] # quiver @@ -943,7 +941,6 @@ def main(args, batches=None): for d in range(args.n_layers): dart_count = int(avg_degree * args.samp_num[d] / avg_degree) args.n_darts.append(dart_count) - print(f"n_darts: {args.n_darts}", flush=True) # for sa-spgemm # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) @@ -1013,6 +1010,7 @@ def main(args, batches=None): # print(f"src_vtxs: {src_vtxs} src_vtxs.size: {src_vtxs.size()} zero_count: {(src_vtxs == 0).sum()}", flush=True) 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) @@ -1057,7 +1055,46 @@ def main(args, batches=None): (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)) - # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") + # 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_row_lens[adj_sample_row_lens.nonzero().squeeze()] + # row_count = adj_sample_row_lens.size(0) + # adj_sample_rows = torch.repeat_interleave( + # torch.arange(0, row_count, + # device=adj_sample_row_lens.device), + # adj_sample_row_lens) + # nnz_col_mask = torch.cuda.BoolTensor(adj_sample_crows[-1].item()).fill_(False) + # rowselect_csr_gpu(adj_sample_skip_cols, adj_sample_crows, nnz_col_mask, \ + # adj_sample_skip_cols.size(0), adj_sample_crows[-1]) + # adj_sample_values = adj_sample_values[~nnz_col_mask] + # adj_sample_cols = adj_sample_cols[~nnz_col_mask] + # else: + # adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + # row_count = adj_row_select_max - adj_row_select_min + # adj_sample_rows = torch.repeat_interleave( + # torch.arange(0, row_count, + # device=adj_sample_row_lens.device), + # adj_sample_row_lens) + + # # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) + # # row_count = args.batch_size * np.prod(args.samp_num[:l], dtype=int) + # col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) + # adj_sample_skip_cols = torch.histc(adj_sample_cols, bins=col_count) + # adj_sample_skip_cols = (adj_sample_skip_cols == 0).nonzero().squeeze(1) + # # print(f"l: {l} adj_sample_cols: {adj_sample_cols} adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + # col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) - \ + # adj_sample_skip_cols.size(0) + # adj_sample_cols = torch.arange(0, adj_sample_values.size(0), device=device, dtype=int) + # # print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + # # print(f"adj_sample_cols: {adj_sample_cols}", flush=True) + # adj_sample_indices = torch.stack((adj_sample_rows, adj_sample_cols)) + # adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ + # adj_sample_values, (row_count, col_count)) + # # (args.batch_size * (args.samp_num ** l), \ + # # args.batch_size * (args.samp_num ** (l + 1)))) + # adj_nnzs.append(adj_sample_values.size(0)) + # # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") if epoch >= 1: model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) @@ -1065,6 +1102,8 @@ def main(args, batches=None): 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)) @@ -1120,8 +1159,14 @@ def main(args, batches=None): 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 + # features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + # features_mask[adj_sample_skip_cols] = False + # features_batch = features_batch[features_mask] else: features_batch = features_loc[src_vtxs] + # features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + # features_mask[adj_sample_skip_cols] = False + # features_batch = features_batch[features_mask] torch.cuda.nvtx.range_pop() # nvtx-selectfeats if epoch >= 1: @@ -1130,13 +1175,16 @@ def main(args, batches=None): # if epoch >= 1: # start_time(start_inner_timer) # torch.cuda.nvtx.range_push("nvtx-fakemats") - # features_batch = torch.cuda.FloatTensor(69740, 128) # SAGEConvfake - # adjs0_rows = torch.randint(0, 9862, (91518,)).cuda().long() # SAGEConvfake - # adjs0_cols = torch.randint(0, 69740, (91518,)).cuda().long() # SAGEConvfake + # 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, 1024, (9065,)).cuda().long() # SAGEConvfake - # adjs1_cols = torch.randint(0, 9862, (9065,)).cuda().long() # 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 From f00273b47a64052ee5f7142aa08bcbd3d8d66b66 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 13 Oct 2023 22:26:38 -0700 Subject: [PATCH 164/178] Added skip columns opt and products loader --- examples/gcn_15d.py | 183 +++++++++++++++++++++++++++----------------- 1 file changed, 112 insertions(+), 71 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 17ada3c..0f38816 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -463,7 +463,32 @@ def main(args, batches=None): data.test_mask = torch.ones(n).long() data = data.to(device) data.y = data.y.to(device) - elif args.dataset.startswith("ogb"): + 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 @@ -772,17 +797,29 @@ def main(args, batches=None): features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ args.replication, args.normalize) print("coalescing", flush=True) - # g_loc = g_loc.coalesce().t_() - 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,:], 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)])) - print("normalizing", flush=True) + if args.dataset == "ogbn-products": + 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: + # g_loc = g_loc.coalesce().t_() + 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,:], 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) + print(f"adj_matrix: {adj_matrix}", flush=True) + 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)])) + print("normalizing", flush=True) g_loc = g_loc.to(device) # g_loc = g_loc.double() # features_loc = inputs @@ -815,10 +852,14 @@ def main(args, batches=None): 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()}") + if args.dataset == "ogbn-papers100M": + train_nid = train_idx + print(f"train_nid.size: {train_nid.size()}") + else: + train_nid = train_idx + test_nid = test_idx + print(f"train_nid.size: {train_nid.size()}") + print(f"test_nid.size: {test_nid.size()}") # # # adj_matrix = None # Uncomment bottom for testing # adj_matrix = adj_matrix.cuda() @@ -1041,60 +1082,60 @@ def main(args, batches=None): adj_matrix_sample_values, \ (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_row_lens[adj_sample_row_lens.nonzero().squeeze()] - # row_count = adj_sample_row_lens.size(0) - # adj_sample_rows = torch.repeat_interleave( - # torch.arange(0, row_count, - # device=adj_sample_row_lens.device), - # adj_sample_row_lens) - # nnz_col_mask = torch.cuda.BoolTensor(adj_sample_crows[-1].item()).fill_(False) - # rowselect_csr_gpu(adj_sample_skip_cols, adj_sample_crows, nnz_col_mask, \ - # adj_sample_skip_cols.size(0), adj_sample_crows[-1]) - # adj_sample_values = adj_sample_values[~nnz_col_mask] - # adj_sample_cols = adj_sample_cols[~nnz_col_mask] - # else: - # adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] - # row_count = adj_row_select_max - adj_row_select_min - # adj_sample_rows = torch.repeat_interleave( - # torch.arange(0, row_count, - # device=adj_sample_row_lens.device), - # adj_sample_row_lens) - - # # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) - # # row_count = args.batch_size * np.prod(args.samp_num[:l], dtype=int) - # col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) - # adj_sample_skip_cols = torch.histc(adj_sample_cols, bins=col_count) - # adj_sample_skip_cols = (adj_sample_skip_cols == 0).nonzero().squeeze(1) - # # print(f"l: {l} adj_sample_cols: {adj_sample_cols} adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) - # col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) - \ - # adj_sample_skip_cols.size(0) - # adj_sample_cols = torch.arange(0, adj_sample_values.size(0), device=device, dtype=int) - # # print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) - # # print(f"adj_sample_cols: {adj_sample_cols}", flush=True) + # adj_sample_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, (row_count, col_count)) + # 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)) - # # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") + 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_row_lens[adj_sample_row_lens.nonzero().squeeze()] + row_count = adj_sample_row_lens.size(0) + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, row_count, + device=adj_sample_row_lens.device), + adj_sample_row_lens) + nnz_col_mask = torch.cuda.BoolTensor(adj_sample_crows[-1].item()).fill_(False) + rowselect_csr_gpu(adj_sample_skip_cols, adj_sample_crows, nnz_col_mask, \ + adj_sample_skip_cols.size(0), adj_sample_crows[-1]) + adj_sample_values = adj_sample_values[~nnz_col_mask] + adj_sample_cols = adj_sample_cols[~nnz_col_mask] + else: + adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] + row_count = adj_row_select_max - adj_row_select_min + adj_sample_rows = torch.repeat_interleave( + torch.arange(0, row_count, + device=adj_sample_row_lens.device), + adj_sample_row_lens) + + # frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) + # row_count = args.batch_size * np.prod(args.samp_num[:l], dtype=int) + col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) + adj_sample_skip_cols = torch.histc(adj_sample_cols, bins=col_count) + adj_sample_skip_cols = (adj_sample_skip_cols == 0).nonzero().squeeze(1) + # print(f"l: {l} adj_sample_cols: {adj_sample_cols} adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) - \ + adj_sample_skip_cols.size(0) + adj_sample_cols = torch.arange(0, adj_sample_values.size(0), device=device, dtype=int) + # print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + # print(f"adj_sample_cols: {adj_sample_cols}", flush=True) + adj_sample_indices = torch.stack((adj_sample_rows, adj_sample_cols)) + adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ + adj_sample_values, (row_count, col_count)) + # (args.batch_size * (args.samp_num ** l), \ + # args.batch_size * (args.samp_num ** (l + 1)))) + adj_nnzs.append(adj_sample_values.size(0)) + # print(f"i: {i} sample.nnz: {adj_matrix_sample._nnz()} sample.size: {adj_matrix_sample.size()} emptyrows: {(adj_sample_row_lens == 0).sum()}") if epoch >= 1: model.timings["extract-inst"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) @@ -1159,14 +1200,14 @@ def main(args, batches=None): 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 - # features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) - # features_mask[adj_sample_skip_cols] = False - # features_batch = features_batch[features_mask] + features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + features_mask[adj_sample_skip_cols] = False + features_batch = features_batch[features_mask] else: features_batch = features_loc[src_vtxs] - # features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) - # features_mask[adj_sample_skip_cols] = False - # features_batch = features_batch[features_mask] + features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + features_mask[adj_sample_skip_cols] = False + features_batch = features_batch[features_mask] torch.cuda.nvtx.range_pop() # nvtx-selectfeats if epoch >= 1: From b209f5fbf87af7a8b3b5dfdccd7c19819c429aad Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 14 Oct 2023 20:43:16 -0700 Subject: [PATCH 165/178] Added replicate_graph as option --- cagnet/samplers/sage.py | 20 +- cagnet/samplers/utils.py | 23 +- examples/gcn_15d.py | 369 +++++++++++++++----------- sparse-extension/sparse_coo_tensor.cu | 9 +- 4 files changed, 239 insertions(+), 182 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 2f9d717..fc4b5d5 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -42,7 +42,7 @@ def stop_time(start_timer, stop_timer, barrier=False, timing_arg=None): 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): + timing_arg, baseline_arg, replicate_graph): global timing global baseline_compare @@ -81,7 +81,8 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total batches_expand = sparse_coo_tensor_gpu(batches_expand_idxs, batches._values(), torch.Size([mb_count * batch_size, node_count_total])) - # batches_expand = batches_expand.to_sparse_csr() + if not replicate_graph: + batches_expand = batches_expand.to_sparse_csr() # adj_matrix = adj_matrix.to_sparse_csr() current_frontier = batches_expand @@ -113,15 +114,18 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # nnz = current_frontier._nnz() // mb_count # nnz = batch_size * (frontier_size ** i) nnz = batch_size * int(np.prod(frontier_sizes[:i], dtype=int)) - # current_frontier = current_frontier.to_sparse_csr() + if not replicate_graph: + current_frontier = current_frontier.to_sparse_csr() # Expand batches matrix if baseline_compare: total_start_timer.record() + print(f"current_frontier: {current_frontier}", flush=True) + print(f"adj_matrix: {adj_matrix}", flush=True) 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) + timing_arg, replicate_graph) # dist.barrier() @@ -129,11 +133,13 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total if p.layout == torch.sparse_csr: p = p.to_sparse_coo() + print(f"p: {p}", flush=True) 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"next_frontier: {next_frontier}", flush=True) # dist.barrier() start_time(start_timer) @@ -171,8 +177,10 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total timing_dict["frontier-row-col-select"].append(stop_time(start_timer, stop_timer)) start_time(start_timer) - # current_frontier_select = current_frontier.col_indices().view(current_frontier._nnz(), 1) - current_frontier_select = current_frontier._indices()[1,:].view(current_frontier._nnz(), 1) + 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 diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index e6f1eb9..4cfcab6 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -11,7 +11,7 @@ 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, \ - reduce_sum_gpu + rearrangel_rows_gpu, reduce_sum_gpu timing = True @@ -574,7 +574,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups if matb.layout == torch.sparse_csr: start_time(start_inner_timer) - rowselect_csr_gpu(nnz_col_ids[j], matb.crow_indices(), nnz_row_mask, \ + 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}") @@ -805,7 +805,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups matc_chunk_crows[0] = 0 matc_chunk_cols = torch.cuda.LongTensor(matc_chunk_crows[-1].item()).fill_(0) - rearrange_rows_gpu(mata_recv_rows, mata_chunk_cols, matc_chunk_crows, + # 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, matb_recv_cols, 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) @@ -1017,7 +1018,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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): + timing_dict, name, timing_arg, replicate_graph): global timing timing = timing_arg @@ -1032,7 +1033,7 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # numerator, adj_matrix, replication, rank, size, # row_groups, col_groups, "prob", sa_masks, # timing_dict, name) - if name == "sage": + 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,:] @@ -1045,8 +1046,8 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication rearrange_rows_gpu(numerator._indices()[0,:], numerator._indices()[1,:], matc_chunk_crows, - adj_matrix.crow_indices().long(), - adj_matrix.col_indices().long(), + adj_matrix.crow_indices(), + adj_matrix.col_indices(), matc_chunk_cols) p_num_rows = torch.repeat_interleave( torch.arange(matc_chunk_crows.size(0) - 1, @@ -1058,7 +1059,11 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication p_num_values = torch.cuda.DoubleTensor(p_num_indices.size(1)).fill_(1.0) else: sa_masks.fill_(False) - # TODO: LADIES code + 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() @@ -1439,7 +1444,6 @@ def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ sa_masks.fill_(0) row_select_mtx = row_select_mtx.to_sparse_csr() # adj_matrix = adj_matrix.to_sparse_csr() - print(f"adj_matrix.layout: {adj_matrix.layout}") sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ row_groups, col_groups, "rowsel", sa_masks, timing_dict) @@ -1639,7 +1643,6 @@ def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ sa_masks.fill_(0) row_select_mtx = row_select_mtx.to_sparse_csr() # adj_matrix = adj_matrix.to_sparse_csr() - print(f"adj_matrix.layout: {adj_matrix.layout}") sampled_indices, sampled_values = dist_saspgemm15D(row_select_mtx, adj_matrix, replication, rank, size, \ row_groups, col_groups, "rowsel", sa_masks, timing_dict) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 0f38816..4c42485 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -290,7 +290,7 @@ def row_normalize(mx): return mx def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, replication, \ - normalize): + 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)) @@ -308,7 +308,10 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep # 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) + 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) @@ -329,14 +332,18 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep # # am_pbyp[i] = scale_elements(adj_matrix, am_pbyp[i], node_count, vtx_indices[i], # # vtx_indices[rank_c], 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)).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) + 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) @@ -346,13 +353,11 @@ def one5d_partition(rank, size, inputs, adj_matrix, data, features, classes, rep input_partitions_fused[-1] = torch.cat(input_partitions[-2:], dim=0) input_partitions = input_partitions_fused - # 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_partitions, input_partitions - return inputs_loc, None, None, input_partitions + return inputs_loc, adj_matrix_loc, am_partitions, input_partitions def one5d_partition_mb(rank, size, batches, replication, mb_count): rank_c = rank // replication @@ -554,83 +559,84 @@ def main(args, batches=None): if rank % args.gpu == 0: if True or args.dataset != "ogbn-papers100M": inputs = inputs.to(torch.device("cpu")) - # adj_matrix = adj_matrix.to(torch.device("cpu")) - # 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) - - print(f"begin conversion", flush=True) - # adj_matrix = adj_matrix.int() - # adj_matrix_crows = torch.IntTensor(inputs.size(0) + 1).fill_(0) - # ones = torch.IntTensor(adj_matrix.size(1)).fill_(1) - # # adj_matrix_rows = torch.histc(adj_matrix[0,:], bins=inputs.size(0)) - # adj_matrix_rows = torch.IntTensor(inputs.size(0)).fill_(0) - # adj_matrix_rows.scatter_add_(0, adj_matrix[0,:].long(), ones) - # adj_matrix_crows[1:] = torch.cumsum(adj_matrix_rows, dim=0).roll(1) - # adj_matrix_crows[0] = 0 - adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") - adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") - print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) - print(f"adj_matrix_cols: {adj_matrix_cols} dtype: {adj_matrix_cols.dtype} device: {adj_matrix_cols.device}", flush=True) - print(f"stop conversion", flush=True) - # adj_matrix_crows = adj_matrix_crows.int().to(device) - # adj_matrix_cols = adj_matrix[1,:].int().to(device) - adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) - g_loc = torch.sparse_csr_tensor(adj_matrix_crows.int(), adj_matrix_cols.int(), \ - adj_matrix_vals, \ - torch.Size([inputs.size(0), inputs.size(0)])) - g_loc = g_loc.to(device) + 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) + + print(f"begin conversion", flush=True) + adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") + adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") + print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) + print(f"adj_matrix_cols: {adj_matrix_cols} dtype: {adj_matrix_cols.dtype} device: {adj_matrix_cols.device}", flush=True) + print(f"stop conversion", flush=True) + # adj_matrix_crows = adj_matrix_crows.int().to(device) + # adj_matrix_cols = adj_matrix[1,:].int().to(device) + adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) + g_loc = torch.sparse_csr_tensor(adj_matrix_crows.int(), adj_matrix_cols.int(), \ + adj_matrix_vals, \ + torch.Size([inputs.size(0), inputs.size(0)])) + g_loc = g_loc.to(device) + else: + adj_matrix = adj_matrix.to(torch.device("cpu")) + features_loc, g_loc, am_partitions, input_partitions = \ + one5d_partition(rank, size, inputs, adj_matrix, data, + inputs, num_classes, args.replication, \ + args.normalize, args.replicate_graph) features_zero_row = input_partitions[0][0].to(device) # feature vector in row 0 print(f"stop copying", flush=True) for i in range(1, args.gpu): print(f"iteration i: {i}", flush=True) - rank_send_row = (rank + i) // args.replication - dst_gpu = rank + i - # g_send = am_partitions[rank_send_row] - g_send = adj_matrix - # 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) - edge_count = g_loc.col_indices().size(0) - - - 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[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 - print(f"g_send_meta: {g_send_meta}", flush=True) - dist.send(g_send_meta, dst=dst_gpu) - # dist.send(g_send._indices(), dst=dst_gpu) - # dist.send(g_send._values(), dst=dst_gpu) - # print(f"i: {i} g_send: {g_send} dtype: {g_send.dtype}", flush=True) - # dist.send(g_send, dst=dst_gpu) - print(f"i: {i} g_loc.crows: {g_loc.crow_indices()} dtype: {g_loc.crow_indices().dtype}", flush=True) - dist.send(g_loc.crow_indices(), dst=dst_gpu) - print(f"i: {i} g_loc.cols: {g_loc.col_indices()} dtype: {g_loc.col_indices().dtype}", flush=True) - dist.send(g_loc.col_indices(), dst=dst_gpu) - print(f"i: {i} features_send.size: {features_send.size()} dtype: {features_send.dtype}", flush=True) - dist.send(features_send, dst=dst_gpu) - print(f"i: {i} features_zero_row.size: {features_zero_row.size()} dtype: {features_zero_row.dtype}", flush=True) - dist.send(features_zero_row, dst=dst_gpu) - print(f"send data.y.size: {data.y.size()} inputs.size: {inputs.size()}", flush=True) - dist.send(data.y.float(), dst=dst_gpu) - # torch.save(g_send_meta, f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_meta.pt") - # torch.save(g_send._indices(), f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_indices.pt") - # torch.save(g_send._values(), f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}g_loc_values.pt") - # torch.save(features_send, f"../../data/ogbn_papers100M/cagnet/p16/{dst_gpu}features_loc.pt") + 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(data.y.float(), dst=dst_gpu) del g_send del features_send @@ -692,64 +698,87 @@ def main(args, batches=None): else: if True or args.dataset != "ogbn-papers100M": 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.IntTensor(2, g_loc_meta[0].item()) - # dist.recv(g_loc_indices, src=src_gpu) - # torch.cuda.synchronize() - # print(f"recv indices", flush=True) + if args.replicate_graph: + g_loc_meta = torch.cuda.LongTensor(6).fill_(0) + dist.recv(g_loc_meta, src=src_gpu) + + # g_loc_indices = torch.cuda.IntTensor(2, g_loc_meta[0].item()) + # dist.recv(g_loc_indices, src=src_gpu) + # torch.cuda.synchronize() + # print(f"recv indices", flush=True) + + g_loc_crows = torch.cuda.IntTensor(g_loc_meta[1].item() + 1) + dist.recv(g_loc_crows, src=src_gpu) + torch.cuda.synchronize() + print(f"recv crows", flush=True) + + g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) + dist.recv(g_loc_cols, src=src_gpu) + torch.cuda.synchronize() + print(f"recv cols", flush=True) + + # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + # dist.recv(g_loc_values, src=src_gpu) + + num_features = g_loc_meta[4].item() + num_classes = g_loc_meta[5].item() + rows_per_proc = g_loc_meta[1].item() // (size // args.replication) + if rank_row < proc_row - 1: + inputs = torch.cuda.FloatTensor(rows_per_proc, num_features) + else: + rows = g_loc_meta[1].item() - rows_per_proc * (proc_row - 1) + inputs = torch.cuda.FloatTensor(rows, num_features) + dist.recv(inputs, src=src_gpu) + torch.cuda.synchronize() + print(f"recv inputs", flush=True) + + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + print(f"data.y.size: {data.y.size()}", flush=True) + dist.recv(data.y, src=src_gpu) + torch.cuda.synchronize() + data.y = data.y.long() + + # g_loc_values = torch.cuda.FloatTensor(g_loc_indices.size(1)).fill_(1.0) + # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + # torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc_values = torch.cuda.FloatTensor(g_loc_cols.size(0)).fill_(1.0) + g_loc_crows = g_loc_crows.cpu() + g_loc_cols = g_loc_cols.cpu() + g_loc_values = g_loc_values.cpu() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc = g_loc.to(device) + edge_count = g_loc_meta[3].item() + else: + g_loc_indices = torch.cuda.LongTensor(2, g_loc_meta[0].item()) + dist.recv(g_loc_indices, src=src_gpu) - g_loc_crows = torch.cuda.IntTensor(g_loc_meta[1].item() + 1) - dist.recv(g_loc_crows, src=src_gpu) - torch.cuda.synchronize() - print(f"recv crows", flush=True) + g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + dist.recv(g_loc_values, src=src_gpu) - g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) - dist.recv(g_loc_cols, src=src_gpu) - torch.cuda.synchronize() - print(f"recv cols", flush=True) + 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) - # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) - # dist.recv(g_loc_values, src=src_gpu) + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + dist.recv(data.y, src=src_gpu) + data.y = data.y.long() - num_features = g_loc_meta[4].item() - num_classes = g_loc_meta[5].item() - rows_per_proc = g_loc_meta[1].item() // (size // args.replication) - if rank_row < proc_row - 1: - inputs = torch.cuda.FloatTensor(rows_per_proc, num_features) - else: - rows = g_loc_meta[1].item() - rows_per_proc * (proc_row - 1) - inputs = torch.cuda.FloatTensor(rows, num_features) - dist.recv(inputs, src=src_gpu) - torch.cuda.synchronize() - print(f"recv inputs", flush=True) - - features_zero_row = torch.cuda.FloatTensor(num_features) - dist.recv(features_zero_row, src=src_gpu) - torch.cuda.synchronize() - print(f"recv features_zero_row", flush=True) - - # data = Data() - data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) - print(f"data.y.size: {data.y.size()}", flush=True) - dist.recv(data.y, src=src_gpu) - torch.cuda.synchronize() - print(f"recv datay", flush=True) - data.y = data.y.long() - print(f"cast datay", flush=True) - - # g_loc_values = torch.cuda.FloatTensor(g_loc_indices.size(1)).fill_(1.0) - # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ - # torch.Size([g_loc_meta[1], g_loc_meta[2]])) - g_loc_values = torch.cuda.FloatTensor(g_loc_cols.size(0)).fill_(1.0) - g_loc_crows = g_loc_crows.cpu() - g_loc_cols = g_loc_cols.cpu() - g_loc_values = g_loc_values.cpu() - g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ - torch.Size([g_loc_meta[1], g_loc_meta[2]])) - g_loc = g_loc.to(device) - edge_count = g_loc_meta[3].item() + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + edge_count = g_loc_meta[3].item() if args.dataset == "ogbn-papers100M": train_idx_len = torch.cuda.LongTensor(1).fill_(0) @@ -795,7 +824,7 @@ def main(args, batches=None): 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.replication, args.normalize, args.replicate_graph) print("coalescing", flush=True) if args.dataset == "ogbn-products": adj_matrix_values = torch.FloatTensor(adj_matrix.size(1)).fill_(1.0) @@ -808,17 +837,19 @@ def main(args, batches=None): 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: - # g_loc = g_loc.coalesce().t_() - 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,:], 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) - print(f"adj_matrix: {adj_matrix}", flush=True) - 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)])) + 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,:], 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) + print(f"adj_matrix: {adj_matrix}", flush=True) + 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_() print("normalizing", flush=True) g_loc = g_loc.to(device) # g_loc = g_loc.double() @@ -888,8 +919,8 @@ def main(args, batches=None): model = model.to(device) # use optimizer - # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) + # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) dur = [] # torch.manual_seed(0) @@ -908,6 +939,7 @@ def main(args, batches=None): g_loc = g_loc.to_sparse_csr() g_loc = g_loc.to(device) + print(f"before adj[1256]: {g_loc.col_indices()[g_loc.crow_indices()[1256]:g_loc.crow_indices()[1257]]}", flush=True) # csr_topo = quiver.CSRTopo(adj_matrix) # quiver_sampler = quiver.pyg.GraphSageSampler(csr_topo, [args.samp_num] * args.n_layers, 0, @@ -985,10 +1017,12 @@ def main(args, batches=None): # for sa-spgemm # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) - # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._nnz()) - # nnz_row_masks.fill_(0) + nnz_row_masks = None + if 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) + # nnz_recv_upperbound = edge_count // (size // args.replication) if epoch >= 1: start_time(start_timer, timing_arg=True) @@ -1003,15 +1037,19 @@ def main(args, batches=None): nnz_row_masks, rank, size, row_groups, col_groups, args.timing) elif args.sample_method == "sage": - # nnz_row_masks.fill_(False) - nnz_row_masks = None + 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, \ - 1, nnz_row_masks, - # args.replication, nnz_row_masks, + rep_pass, nnz_row_masks, rank, size, row_groups, - col_groups, args.timing, args.baseline) + col_groups, args.timing, args.baseline, + args.replicate_graph) if epoch >= 1: model.timings["sample"].append(stop_time(start_timer, stop_timer, timing_arg=True)) @@ -1258,6 +1296,9 @@ def main(args, batches=None): # adjs = [adj.to(rank) for adj in adjs] # features_batch = features_loc[n_id] optimizer.zero_grad() + print(f"adjs: {adjs}", flush=True) + print(f"features_batch: {features_batch}", flush=True) + print(f"src_vtxs: {src_vtxs}", flush=True) logits = model(adjs, features_batch, epoch) torch.cuda.nvtx.range_pop() # nvtx-fwd if epoch >= 1: @@ -1288,6 +1329,8 @@ def main(args, batches=None): torch.cuda.nvtx.range_push("nvtx-optstep") optimizer.step() torch.cuda.nvtx.range_pop() # nvtx-optstep + # del adjs + # del features_batch if epoch >= 1: model.timings["train"].append(stop_time(start_timer, stop_timer, timing_arg=True)) @@ -1417,6 +1460,8 @@ def main(args, batches=None): 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) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index f62cdeb..f84f81f 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2691,8 +2691,8 @@ void rearrangel_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_col CHECK_ERROR("rearrange_rows error") } -__global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, long *matb_crows, - long *matb_cols, int *matc_cols, long nnz_count) { +__global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, int *matb_crows, + int *matb_cols, int *matc_cols, long nnz_count) { int id = blockIdx.x * blockDim.x + threadIdx.x; int stride = blockDim.x * gridDim.x; @@ -2726,8 +2726,8 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols RearrangeRows<<>>(mata_rows.data(), mata_cols.data(), matc_crows.data(), - matb_crows.data(), - matb_cols.data(), + matb_crows.data(), + matb_cols.data(), matc_cols.data(), mata_rows.sizes()[0]); @@ -2843,5 +2843,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From 70e7d3c50a0f81e36a2bf5db5f7442e878856401 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 16 Oct 2023 12:52:50 -0700 Subject: [PATCH 166/178] Some memory saving optimizations --- cagnet/samplers/sage.py | 8 +-- cagnet/samplers/utils.py | 41 +++++++----- examples/gcn_15d.py | 134 +++++++++++++++++++++++++++++---------- 3 files changed, 129 insertions(+), 54 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index fc4b5d5..892c917 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -120,8 +120,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # Expand batches matrix if baseline_compare: total_start_timer.record() - print(f"current_frontier: {current_frontier}", flush=True) - print(f"adj_matrix: {adj_matrix}", flush=True) 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", @@ -133,13 +131,11 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total if p.layout == torch.sparse_csr: p = p.to_sparse_coo() - print(f"p: {p}", flush=True) 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"next_frontier: {next_frontier}", flush=True) # dist.barrier() start_time(start_timer) @@ -210,6 +206,10 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total frontiers[i + 1] = next_frontier_select else: del next_frontier_select + del p + + if i > 0: + del current_frontier_select current_frontier = next_frontier timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer)) adj_matrices[i] = adj_matrix_sample.to_sparse_csr() diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 4cfcab6..15da1ee 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -570,7 +570,8 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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 = 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) @@ -622,11 +623,13 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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.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.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: @@ -707,25 +710,27 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups req.wait() torch.cuda.synchronize() - matb_recv_lengths = matb_recv_lengths.long() - matb_recv_cols = matb_recv_cols.long() - # 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, dim=0) + matb_recv_crows[1:] = torch.cumsum(matb_rows_sum, dtypes=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.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[1:] = torch.cumsum(matb_rows_sum, dim=0) + # 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) @@ -799,14 +804,16 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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.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] 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) - # rearrange_rows_gpu(mata_recv_rows, mata_chunk_cols, matc_chunk_crows, - rearrangel_rows_gpu(mata_recv_rows, mata_chunk_cols, matc_chunk_crows, + # 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, matb_recv_cols, 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) @@ -829,7 +836,6 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # matc_chunk_cols = matc_chunk.col_indices() # matc_chunk_values = matc_chunk.values() matc_row_lens = matc_chunk_crows[1:] - matc_chunk_crows[:-1] - lt_zero_mask = matc_row_lens < 0 matc_chunk_rows = torch.repeat_interleave( torch.arange(0, mata.size(0), device=mata.device), matc_row_lens) @@ -1058,7 +1064,8 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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) else: - sa_masks.fill_(False) + 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, diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 4c42485..bb217db 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -85,6 +85,7 @@ def forward(self, graphs, inputs, epoch): 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 @@ -445,7 +446,7 @@ def main(args, batches=None): 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 * 2048] = 1 + data.train_mask[:args.batch_size * 1024] = 1 data.test_mask = torch.ones(n).long() adj_matrix = edge_index.t_() data = data.to(device) @@ -519,18 +520,43 @@ def main(args, batches=None): # split_idx = dataset.get_idx_split() # train_idx = split_idx['train'].to(device) # test_idx = split_idx['test'].to(device) - adj_matrix = None - print(f"before loading", flush=True) - data.y = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/label/label.pt") - data.y = data.y.squeeze().to(device) - print(f"data.y: {data.y} data.y.dtype {data.y.dtype}", flush=True) - train_idx = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/index/train_idx.pt") - train_idx = train_idx.to(device) - print(f"train_idx: {train_idx} train_idx.dtype {train_idx.dtype}", flush=True) - inputs = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") - print(f"inputs: {inputs} inputs.dtype {inputs.dtype}", flush=True) - num_features = inputs.size(1) - num_classes = 172 + if not args.replicate_graph: + import ogb + from ogb.nodeproppred import PygNodePropPredDataset # only import if necessary, takes a long time + dataset = PygNodePropPredDataset(name=args.dataset, root="/global/u1/a/alokt/data") + data = dataset[0] + + # split_idx = dataset.get_idx_split() + # # train_loader = DataLoader(dataset[split_idx["train"]], batch_size=32, shuffle=True) + # # valid_loader = DataLoader(dataset[split_idx["valid"]], batch_size=32, shuffle=False) + # # test_loader = DataLoader(dataset[split_idx["test"]], batch_size=32, shuffle=False) + # # data = data.to(device) + # # data.x.requires_grad = True + # # inputs = data.x.to(device) + inputs = data.x + data.y = data.y.squeeze().to(device) + # # inputs.requires_grad = True + # # data.y = data.y.to(device) + num_features = dataset.num_features + edge_index = data.edge_index + num_classes = dataset.num_classes + adj_matrix = edge_index + split_idx = dataset.get_idx_split() + train_idx = split_idx['train'].to(device) + test_idx = split_idx['test'].to(device) + else: + adj_matrix = None + print(f"before loading", flush=True) + data.y = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/label/label.pt") + data.y = data.y.squeeze().to(device) + print(f"data.y: {data.y} data.y.dtype {data.y.dtype}", flush=True) + train_idx = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/index/train_idx.pt") + train_idx = train_idx.to(device) + print(f"train_idx: {train_idx} train_idx.dtype {train_idx.dtype}", flush=True) + inputs = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") + print(f"inputs: {inputs} inputs.dtype {inputs.dtype}", flush=True) + num_features = inputs.size(1) + num_classes = 172 # print(f"adj_matrix.size(): {adj_matrix.size()}") # print(f"adj_matrix: {adj_matrix}") @@ -636,6 +662,7 @@ def main(args, batches=None): 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 @@ -648,12 +675,20 @@ def main(args, batches=None): # dist.send(test_idx_len, dst=dst_gpu) dist.send(train_idx, dst=dst_gpu) # dist.send(test_idx, dst=dst_gpu) - print("coalescing", flush=True) - # g_loc = g_loc.coalesce().t_() - print("normalizing", flush=True) + 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) features_loc = features_loc.to(device) - edge_count = g_loc.col_indices().size(0) else: print(f"loading", flush=True) g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", @@ -677,7 +712,11 @@ def main(args, batches=None): g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, torch.Size([g_loc_meta[1], g_loc_meta[2]])) print("normalizing", flush=True) - edge_count = g_loc.col_indices().size(0) + if args.replicate_graph: + edge_count = g_loc.col_indices().size(0) + else: + edge_count = g_loc._indices().size(1) + g_loc = g_loc.coalesce().t_() # features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ # args.replication, args.normalize) @@ -756,6 +795,10 @@ def main(args, batches=None): g_loc = g_loc.to(device) edge_count = g_loc_meta[3].item() else: + src_gpu = rank - (rank % args.gpu) + g_loc_meta = torch.cuda.LongTensor(6).fill_(0) + dist.recv(g_loc_meta, src=src_gpu) + g_loc_indices = torch.cuda.LongTensor(2, g_loc_meta[0].item()) dist.recv(g_loc_indices, src=src_gpu) @@ -771,6 +814,11 @@ def main(args, batches=None): inputs = torch.cuda.FloatTensor(rows, num_features) dist.recv(inputs, src=src_gpu) + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + # data = Data() data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) dist.recv(data.y, src=src_gpu) @@ -778,6 +826,12 @@ def main(args, batches=None): g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ torch.Size([g_loc_meta[1], g_loc_meta[2]])) + 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) edge_count = g_loc_meta[3].item() if args.dataset == "ogbn-papers100M": @@ -826,7 +880,7 @@ def main(args, batches=None): 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": + 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)])) @@ -840,16 +894,19 @@ def main(args, batches=None): 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,:], ones) + 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) - print(f"adj_matrix: {adj_matrix}", flush=True) 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() @@ -937,10 +994,12 @@ def main(args, batches=None): print(f"rank_n_bulkmb: {rank_n_bulkmb}") # g_loc = g_loc.cpu() - g_loc = g_loc.to_sparse_csr() + # 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) - print(f"before adj[1256]: {g_loc.col_indices()[g_loc.crow_indices()[1256]:g_loc.crow_indices()[1257]]}", flush=True) - # csr_topo = quiver.CSRTopo(adj_matrix) # quiver_sampler = quiver.pyg.GraphSageSampler(csr_topo, [args.samp_num] * args.n_layers, 0, # mode='UVA') @@ -1018,7 +1077,7 @@ def main(args, batches=None): # for sa-spgemm # nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._indices().size(1)) nnz_row_masks = None - if not args.replicate_graph: + if False and not args.replicate_graph: nnz_row_masks = torch.cuda.BoolTensor((size // args.replication) * g_loc._nnz()) nnz_row_masks.fill_(0) @@ -1086,7 +1145,6 @@ def main(args, batches=None): src_select_max = (i + 1) * args.batch_size * np.prod(args.samp_num[:-1], dtype=int) src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) - # print(f"src_vtxs: {src_vtxs} src_vtxs.size: {src_vtxs.size()} zero_count: {(src_vtxs == 0).sum()}", flush=True) adjs = [None] * args.n_layers adj_sample_skip_cols = None @@ -1136,8 +1194,10 @@ def main(args, batches=None): # 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_row_lens[adj_sample_row_lens.nonzero().squeeze()] + # 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, @@ -1151,6 +1211,7 @@ def main(args, batches=None): 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"row_count: {row_count} min: {adj_row_select_min} max: {adj_row_select_max}", flush=True) adj_sample_rows = torch.repeat_interleave( torch.arange(0, row_count, device=adj_sample_row_lens.device), @@ -1296,9 +1357,6 @@ def main(args, batches=None): # adjs = [adj.to(rank) for adj in adjs] # features_batch = features_loc[n_id] optimizer.zero_grad() - print(f"adjs: {adjs}", flush=True) - print(f"features_batch: {features_batch}", flush=True) - print(f"src_vtxs: {src_vtxs}", flush=True) logits = model(adjs, features_batch, epoch) torch.cuda.nvtx.range_pop() # nvtx-fwd if epoch >= 1: @@ -1308,6 +1366,7 @@ def main(args, batches=None): start_time(start_inner_timer) torch.cuda.nvtx.range_push("nvtx-loss") # loss = F.nll_loss(logits, data.y[batch_vtxs].long()) # GCNConv + print(f"logits.size: {logits.size()}", flush=True) loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs].long()) # SAGEConv # loss = F.nll_loss(logits[:args.batch_size], data.y[seeds].long()) # quiver @@ -1328,9 +1387,11 @@ def main(args, batches=None): W.grad /= size torch.cuda.nvtx.range_push("nvtx-optstep") optimizer.step() + + for i in range(len(adjs)): + del adjs[0] + del adjs torch.cuda.nvtx.range_pop() # nvtx-optstep - # del adjs - # del features_batch if epoch >= 1: model.timings["train"].append(stop_time(start_timer, stop_timer, timing_arg=True)) @@ -1339,7 +1400,14 @@ def main(args, batches=None): last_batch = False args.n_bulkmb = tmp_bulkmb args.batch_size = tmp_batch_size + torch.cuda.nvtx.range_pop() # nvtx-training + for i in range(len(frontiers_bulk)): + del frontiers_bulk[0] + del frontiers_bulk + for i in range(len(adj_matrices_bulk)): + del adj_matrices_bulk[0] + del adj_matrices_bulk torch.cuda.nvtx.range_pop() # nvtx-bulk dist.barrier() From 67cc4d78c791aa130418fb6aed5a234d0efc46f1 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 17 Oct 2023 18:29:35 -0700 Subject: [PATCH 167/178] Cora LADIES p4c2 --- cagnet/samplers/ladies.py | 13 +- cagnet/samplers/utils.py | 84 ++-- examples/gcn_15d.py | 665 +++++++++++++++----------- sparse-extension/sparse_coo_tensor.cu | 4 +- 4 files changed, 441 insertions(+), 325 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 076fe9a..60e94ff 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -25,13 +25,14 @@ def stop_time(start_timer, stop_timer, barrier=False): 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): + 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) @@ -57,9 +58,12 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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.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) + 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") @@ -74,7 +78,8 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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 + adj_matrices[i] = adj_matrix_sample.to_sparse_csr() current_frontier = next_frontier diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index 15da1ee..c2c7ad9 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -614,11 +614,14 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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.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.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) @@ -718,7 +721,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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, dtypes=torch.int32, dim=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_cols, matb_recv_lengths) @@ -1156,6 +1159,8 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # 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: @@ -1181,36 +1186,49 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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() - 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()) - # n_darts_col * p.size(0), p._nnz()) - # n_darts * mb_count, p._nnz()) - timing_dict["sample-filter-darts"].append(stop_time(start_timer, stop_timer)) + 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) - # dist.all_reduce(dart_hits_count, group=row_groups[rank_c]) + 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() @@ -1338,8 +1356,9 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, timing_dict["sample-dart-selection"].append(stop_time(select_start_timer, select_stop_timer)) start_time(start_timer) - # 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]) + 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() @@ -1647,7 +1666,8 @@ def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ # 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])) - sa_masks.fill_(0) + 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, \ diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index bb217db..49e6607 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -85,7 +85,7 @@ def forward(self, graphs, inputs, epoch): 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) + # 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 @@ -194,6 +194,116 @@ def evaluate(self, graph, features, test_idx, labels): # return features +class LADIES(nn.Module): + def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, partitioning, replication, + device, group=None, row_groups=None, col_groups=None): + super(LADIES, self).__init__() + self.layers = nn.ModuleList() + self.n_layers = n_layers + self.aggr = aggr + self.rank = rank + self.size = size + self.group = group + self.row_groups = row_groups + self.col_groups = col_groups + self.device = device + self.partitioning = partitioning + self.replication = replication + self.timings = dict() + + self.timings["total"] = [] + self.timings["sample"] = [] + self.timings["extract"] = [] + self.timings["extract-select"] = [] + self.timings["extract-inst"] = [] + self.timings["extract-coalesce"] = [] + self.timings["train"] = [] + self.timings["selectfeats"] = [] + self.timings["fwd"] = [] + self.timings["bwd"] = [] + self.timings["loss"] = [] + self.timings["fakemats"] = [] + + self.timings["precomp"] = [] + self.timings["spmm"] = [] + self.timings["gemm_i"] = [] + self.timings["gemm_w"] = [] + self.timings["aggr"] = [] + + # # input layer + # self.layers.append(GCNConv(in_feats, n_hidden, self.partitioning, self.device)) + # # hidden layers + # for i in range(n_layers - 2): + # self.layers.append(GCNConv(n_hidden, n_hidden, self.partitioning, self.device)) + # # output layer + # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) + self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) + + def forward(self, graphs, inputs, epoch): + h = inputs + for l, layer in enumerate(self.layers): + # graphs[l].t_() + # edge_index = graphs[l]._indices() + h = layer(self, graphs[l], h, epoch) # GCNConv + if l != len(self.layers) - 1: + # h = CAGF.relu(h, self.partitioning) + h = F.relu(h) + + # h = CAGF.log_softmax(self, h, self.partitioning, dim=1) + h = F.log_softmax(h, dim=1) + return h + + @torch.no_grad() + def evaluate(self, graph, features, test_idx, labels): + # subgraph_loader = NeighborSampler(graph, node_idx=None, + # sizes=[-1], batch_size=2048, + # shuffle=False, num_workers=6) + + # for i in range(self.n_layers): + # xs = [] + # for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + # x = features[n_id].to(self.device) + # # edge_index, _, size = adj + # # x = features[n_id] + # # x_target = x[:size[1]] + # # x = self.convs[i]((x, x_target), edge_index) + # x = self.layers[i](x, edge_index) + # if i != self.n_layers - 1: + # x = F.relu(x) + # # xs.append(x) + # xs.append(x[:batch_size]) + + # features = torch.cat(xs, dim=0) + + # return features + + subgraph_loader = NeighborSampler(graph, node_idx=None, + sizes=[-1], batch_size=2048, + # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, + # sizes=[-1, -1], batch_size=512, + shuffle=False) + non_eval_timings = copy.deepcopy(self.timings) + for l, layer in enumerate(self.layers): + hs = [] + for batch_size, n_id, adj in subgraph_loader: + # edge_index, _, size = adj.to(self.device) + edge_index, _, size = adj + adj_batch = torch.sparse_coo_tensor(edge_index, + torch.FloatTensor(edge_index.size(1)).fill_(1.0), + size) + adj_batch = adj_batch.t().coalesce() + h = features[n_id] + + # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv + h = self.layers[l](h, edge_index) + if l != len(self.layers) - 1: + h = CAGF.relu(h, self.partitioning) + # hs.append(h) # GCNConv + hs.append(h[:batch_size]) # SAGEConv + features = torch.cat(hs, dim=0) + return features + def get_proc_groups(rank, size, replication): rank_c = rank // replication @@ -583,140 +693,120 @@ def main(args, batches=None): torch.cuda.synchronize() if args.dataset == "ogbn-papers100M" or "Protein" in args.dataset: if rank % args.gpu == 0: - if True or args.dataset != "ogbn-papers100M": - 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) + 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) - print(f"begin conversion", flush=True) + if args.dataset == "ogbn-papers100M": + print(f"begin copying", flush=True) adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) print(f"adj_matrix_cols: {adj_matrix_cols} dtype: {adj_matrix_cols.dtype} device: {adj_matrix_cols.device}", flush=True) - print(f"stop conversion", flush=True) # adj_matrix_crows = adj_matrix_crows.int().to(device) # adj_matrix_cols = adj_matrix[1,:].int().to(device) adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) g_loc = torch.sparse_csr_tensor(adj_matrix_crows.int(), adj_matrix_cols.int(), \ adj_matrix_vals, \ torch.Size([inputs.size(0), inputs.size(0)])) - g_loc = g_loc.to(device) else: - adj_matrix = adj_matrix.to(torch.device("cpu")) - features_loc, g_loc, am_partitions, input_partitions = \ - one5d_partition(rank, size, inputs, adj_matrix, data, - inputs, num_classes, args.replication, \ - args.normalize, args.replicate_graph) - features_zero_row = input_partitions[0][0].to(device) # feature vector in row 0 - print(f"stop copying", flush=True) - for i in range(1, args.gpu): - print(f"iteration i: {i}", flush=True) - if args.replicate_graph: - rank_send_row = (rank + i) // args.replication - dst_gpu = rank + i - g_send = adj_matrix - print("normalizing", flush=True) - - features_send = input_partitions[rank_send_row].to(device) - edge_count = g_loc.col_indices().size(0) - - - g_send_meta = torch.cuda.LongTensor(6) - g_send_meta[0] = g_loc.col_indices().size(0) - g_send_meta[1] = inputs.size(0) - g_send_meta[2] = inputs.size(0) - g_send_meta[3] = g_loc.col_indices().size(0) - g_send_meta[4] = num_features - g_send_meta[5] = num_classes - dist.send(g_send_meta, dst=dst_gpu) - dist.send(g_loc.crow_indices(), dst=dst_gpu) - dist.send(g_loc.col_indices(), dst=dst_gpu) - dist.send(features_send, dst=dst_gpu) - dist.send(features_zero_row, dst=dst_gpu) - dist.send(data.y.float(), dst=dst_gpu) - else: - rank_send_row = (rank + i) // args.replication - dst_gpu = rank + i - g_send = am_partitions[rank_send_row] - print("coalescing", flush=True) - g_send = g_send.coalesce().t_() - print("normalizing", flush=True) - g_send = g_send.to(device) - g_send = g_send.double() - features_send = input_partitions[rank_send_row].to(device) - edge_count = adj_matrix.size(1) - - - g_send_meta = torch.cuda.LongTensor(6) - g_send_meta[0] = g_send._nnz() - g_send_meta[1] = g_send.size(0) - g_send_meta[2] = g_send.size(1) - g_send_meta[3] = adj_matrix.size(1) - g_send_meta[4] = num_features - g_send_meta[5] = num_classes - dist.send(g_send_meta, dst=dst_gpu) - dist.send(g_send._indices(), dst=dst_gpu) - dist.send(g_send._values(), dst=dst_gpu) - dist.send(features_send, dst=dst_gpu) - dist.send(features_zero_row, dst=dst_gpu) - dist.send(data.y.float(), dst=dst_gpu) - - del g_send - del features_send - - if args.dataset == "ogbn-papers100M": - train_idx_len = torch.cuda.LongTensor(1).fill_(train_idx.size(0)) - # test_idx_len = torch.cuda.LongTensor(1).fill_(test_idx.size(0)) - dist.send(train_idx_len, dst=dst_gpu) - # dist.send(test_idx_len, dst=dst_gpu) - dist.send(train_idx, dst=dst_gpu) - # dist.send(test_idx, dst=dst_gpu) - if args.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() + 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, g_loc.values(), g_loc.size()) - print("host2device", flush=True) + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, adj_matrix_vals, g_loc.size()) g_loc = g_loc.to(device) - features_loc = features_loc.to(device) else: - print(f"loading", flush=True) - g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", - map_location=torch.device("cpu")) - g_loc_indices = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_indices.pt", - map_location=torch.device("cpu")) - g_loc_values = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_values.pt", - map_location=torch.device("cpu")) - features_loc = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}features_loc.pt", - map_location=torch.device("cpu")) - - g_loc_meta = g_loc_meta.to(device) - g_loc_indices = g_loc_indices.to(device) - g_loc_values = g_loc_values.to(device) - features_loc = features_loc.to(device) - edge_count = g_loc_meta[3].item() - num_features = g_loc_meta[4].item() - num_classes = g_loc_meta[5].item() - # g_loc = sparse_coo_tensor_gpu(g_loc_indices, g_loc_values, - # torch.Size([g_loc_meta[1], g_loc_meta[2]])) - g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, - torch.Size([g_loc_meta[1], g_loc_meta[2]])) - print("normalizing", flush=True) + 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: - edge_count = g_loc._indices().size(1) - g_loc = g_loc.coalesce().t_() + 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.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) + 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) @@ -735,143 +825,114 @@ def main(args, batches=None): # del adj_matrix # Comment when testing # del inputs else: - if True or args.dataset != "ogbn-papers100M": - src_gpu = rank - (rank % args.gpu) - if args.replicate_graph: - g_loc_meta = torch.cuda.LongTensor(6).fill_(0) - dist.recv(g_loc_meta, src=src_gpu) - - # g_loc_indices = torch.cuda.IntTensor(2, g_loc_meta[0].item()) - # dist.recv(g_loc_indices, src=src_gpu) - # torch.cuda.synchronize() - # print(f"recv indices", flush=True) - - g_loc_crows = torch.cuda.IntTensor(g_loc_meta[1].item() + 1) - dist.recv(g_loc_crows, src=src_gpu) - torch.cuda.synchronize() - print(f"recv crows", flush=True) - - g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) - dist.recv(g_loc_cols, src=src_gpu) - torch.cuda.synchronize() - print(f"recv cols", flush=True) - - # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) - # dist.recv(g_loc_values, src=src_gpu) - - num_features = g_loc_meta[4].item() - num_classes = g_loc_meta[5].item() - rows_per_proc = g_loc_meta[1].item() // (size // args.replication) - if rank_row < proc_row - 1: - inputs = torch.cuda.FloatTensor(rows_per_proc, num_features) - else: - rows = g_loc_meta[1].item() - rows_per_proc * (proc_row - 1) - inputs = torch.cuda.FloatTensor(rows, num_features) - dist.recv(inputs, src=src_gpu) - torch.cuda.synchronize() - print(f"recv inputs", flush=True) - - features_zero_row = torch.cuda.FloatTensor(num_features) - dist.recv(features_zero_row, src=src_gpu) - torch.cuda.synchronize() - print(f"recv features_zero_row", flush=True) - - # data = Data() - data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) - print(f"data.y.size: {data.y.size()}", flush=True) - dist.recv(data.y, src=src_gpu) - torch.cuda.synchronize() - data.y = data.y.long() - - # g_loc_values = torch.cuda.FloatTensor(g_loc_indices.size(1)).fill_(1.0) - # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ - # torch.Size([g_loc_meta[1], g_loc_meta[2]])) - g_loc_values = torch.cuda.FloatTensor(g_loc_cols.size(0)).fill_(1.0) - g_loc_crows = g_loc_crows.cpu() - g_loc_cols = g_loc_cols.cpu() - g_loc_values = g_loc_values.cpu() - g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ - torch.Size([g_loc_meta[1], g_loc_meta[2]])) - g_loc = g_loc.to(device) - edge_count = g_loc_meta[3].item() - else: - src_gpu = rank - (rank % args.gpu) - g_loc_meta = torch.cuda.LongTensor(6).fill_(0) - dist.recv(g_loc_meta, src=src_gpu) + 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.LongTensor(2, g_loc_meta[0].item()) - dist.recv(g_loc_indices, 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_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) - dist.recv(g_loc_values, src=src_gpu) + g_loc_crows = torch.cuda.IntTensor(g_loc_meta[1].item() + 1) + dist.recv(g_loc_crows, src=src_gpu) + torch.cuda.synchronize() + print(f"recv crows", flush=True) - num_features = g_loc_meta[4].item() - num_classes = g_loc_meta[5].item() - if rank_row < proc_row - 1: - inputs = torch.cuda.FloatTensor(g_loc_meta[1].item(), num_features) - else: - rows = g_loc_meta[1].item() + (g_loc_meta[2].item() - proc_row * g_loc_meta[1].item()) - inputs = torch.cuda.FloatTensor(rows, num_features) - dist.recv(inputs, src=src_gpu) - - features_zero_row = torch.cuda.FloatTensor(num_features) - dist.recv(features_zero_row, src=src_gpu) - torch.cuda.synchronize() - print(f"recv features_zero_row", flush=True) - - # data = Data() - data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) - dist.recv(data.y, src=src_gpu) - data.y = data.y.long() - - g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ - torch.Size([g_loc_meta[1], g_loc_meta[2]])) - 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) - edge_count = g_loc_meta[3].item() + g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) + dist.recv(g_loc_cols, src=src_gpu) + torch.cuda.synchronize() + print(f"recv cols", flush=True) - 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) - else: - print(f"loading", flush=True) - # Load data to CPU instead of GPU bc multiple processes spawn for some reason, - # and pytorch reserves less memory causing OOMs - g_loc_meta = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_meta.pt", - map_location=torch.device("cpu")) - g_loc_indices = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_indices.pt", - map_location=torch.device("cpu")) - g_loc_values = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}g_loc_values.pt", - map_location=torch.device("cpu")) - inputs = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/{rank}features_loc.pt", - map_location=torch.device("cpu")) - - g_loc_meta = g_loc_meta.to(device) - g_loc_indices = g_loc_indices.to(device) - g_loc_values = g_loc_values.to(device) - inputs = inputs.to(device) + # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + # dist.recv(g_loc_values, src=src_gpu) + + num_features = g_loc_meta[4].item() + num_classes = g_loc_meta[5].item() + rows_per_proc = g_loc_meta[1].item() // (size // args.replication) + if rank_row < proc_row - 1: + inputs = torch.cuda.FloatTensor(rows_per_proc, num_features) + else: + rows = g_loc_meta[1].item() - rows_per_proc * (proc_row - 1) + inputs = torch.cuda.FloatTensor(rows, num_features) + dist.recv(inputs, src=src_gpu) + torch.cuda.synchronize() + print(f"recv inputs", flush=True) + + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + print(f"data.y.size: {data.y.size()}", flush=True) + dist.recv(data.y, src=src_gpu) + torch.cuda.synchronize() + data.y = data.y.long() + + # g_loc_values = torch.cuda.FloatTensor(g_loc_indices.size(1)).fill_(1.0) + # g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + # torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc_values = torch.cuda.FloatTensor(g_loc_cols.size(0)).fill_(1.0) + g_loc_crows = g_loc_crows.cpu() + g_loc_cols = g_loc_cols.cpu() + g_loc_values = g_loc_values.cpu() + g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + g_loc = g_loc.to(device) edge_count = g_loc_meta[3].item() + else: + src_gpu = rank - (rank % args.gpu) + g_loc_meta = torch.cuda.LongTensor(6).fill_(0) + dist.recv(g_loc_meta, src=src_gpu) + + g_loc_indices = torch.cuda.LongTensor(2, g_loc_meta[0].item()) + dist.recv(g_loc_indices, src=src_gpu) + + g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) + dist.recv(g_loc_values, src=src_gpu) + num_features = g_loc_meta[4].item() num_classes = g_loc_meta[5].item() - g_loc = sparse_coo_tensor_gpu(g_loc_indices, g_loc_values, - torch.Size([g_loc_meta[1], g_loc_meta[2]])) - data.y = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/datay.pt") - train_idx = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/train_idx.pt") - test_idx = torch.load(f"../../data/ogbn_papers100M/cagnet/p16/test_idx.pt") - data.y = data.y.to(device) - train_idx = train_idx.to(device) - test_idx = test_idx.to(device) + if rank_row < proc_row - 1: + inputs = torch.cuda.FloatTensor(g_loc_meta[1].item(), num_features) + else: + rows = g_loc_meta[1].item() + (g_loc_meta[2].item() - proc_row * g_loc_meta[1].item()) + inputs = torch.cuda.FloatTensor(rows, num_features) + dist.recv(inputs, src=src_gpu) + + features_zero_row = torch.cuda.FloatTensor(num_features) + dist.recv(features_zero_row, src=src_gpu) + torch.cuda.synchronize() + print(f"recv features_zero_row", flush=True) + + # data = Data() + data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) + dist.recv(data.y, src=src_gpu) + data.y = data.y.long() + + g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ + torch.Size([g_loc_meta[1], g_loc_meta[2]])) + 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) + 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: @@ -961,18 +1022,32 @@ def main(args, batches=None): # create GCN model # torch.manual_seed(0) - 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) + 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 @@ -1050,7 +1125,10 @@ def main(args, batches=None): args.batch_size = batch_count - b torch.cuda.nvtx.range_push("nvtx-partition-batches") batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) - batches_loc = one5d_partition_mb(rank, size, batches, 1, args.n_bulkmb) + 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)) @@ -1094,7 +1172,9 @@ def main(args, batches=None): args.n_layers, args.n_darts, \ args.semibulk, args.replication, \ nnz_row_masks, rank, size, - row_groups, col_groups, args.timing) + 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) @@ -1126,26 +1206,20 @@ def main(args, batches=None): # batch_vtxs = frontiers[0][i].view(-1) torch.cuda.nvtx.range_push("nvtx-extracting") - batch_select_size = args.batch_size - batch_select_min = i * args.batch_size - batch_select_max = (i + 1) * args.batch_size + 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 - # rank_n_bulkmb_row = rank_n_bulkmb // args.replication - # if rank_col == args.replication - 1: - # rank_n_bulkmb_row = rank_n_bulkmb - rank_n_bulkmb_row * (args.replication - 1) + batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) - # batch_select_min += batch_select_size * rank_col * rank_n_bulkmb - # batch_select_max += batch_select_size * rank_col * rank_n_bulkmb - batch_vtxs = frontiers_bulk[0][batch_select_min:batch_select_max,:].view(-1) - - # src_vtxs = frontiers[-1][i].view(-1) - # src_select_min = i * args.batch_size * (args.samp_num ** (args.n_layers - 1)) - # src_select_max = (i + 1) * args.batch_size * (args.samp_num ** (args.n_layers - 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[:-1], dtype=int) + src_select_max = (i + 1) * args.batch_size * np.prod(args.samp_num[:-1], dtype=int) + src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) + elif args.sample_method == "ladies": + batch_vtxs = frontiers_bulk[0][i] + src_vtxs = frontiers_bulk[-1][i] - src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) - adjs = [None] * args.n_layers adj_sample_skip_cols = None for l in range(args.n_layers): @@ -1174,9 +1248,23 @@ def main(args, batches=None): if epoch >= 1: start_time(start_inner_timer, timing_arg=True) if args.sample_method == "ladies": - adj_matrix_sample = torch.sparse_coo_tensor(adj_matrix_sample_indices, \ - adj_matrix_sample_values, \ - (args.batch_size, args.samp_num + args.batch_size)) + 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 @@ -1211,7 +1299,6 @@ def main(args, batches=None): 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"row_count: {row_count} min: {adj_row_select_min} max: {adj_row_select_max}", flush=True) adj_sample_rows = torch.repeat_interleave( torch.arange(0, row_count, device=adj_sample_row_lens.device), @@ -1251,6 +1338,7 @@ def main(args, batches=None): if epoch >= 1: start_time(start_inner_timer, timing_arg=True) torch.cuda.nvtx.range_push("nvtx-selectfeats") + print(f"adjs: {adjs}", flush=True) if size > 1: src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) @@ -1299,14 +1387,16 @@ def main(args, batches=None): 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 - features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) - features_mask[adj_sample_skip_cols] = False - features_batch = features_batch[features_mask] + if adj_sample_skip_cols is not None: + features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + features_mask[adj_sample_skip_cols] = False + features_batch = features_batch[features_mask] else: features_batch = features_loc[src_vtxs] - features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) - features_mask[adj_sample_skip_cols] = False - features_batch = features_batch[features_mask] + if adj_sample_skip_cols is not None: + features_mask = torch.cuda.BoolTensor(features_batch.size(0)).fill_(True) + features_mask[adj_sample_skip_cols] = False + features_batch = features_batch[features_mask] torch.cuda.nvtx.range_pop() # nvtx-selectfeats if epoch >= 1: @@ -1366,7 +1456,6 @@ def main(args, batches=None): start_time(start_inner_timer) torch.cuda.nvtx.range_push("nvtx-loss") # loss = F.nll_loss(logits, data.y[batch_vtxs].long()) # GCNConv - print(f"logits.size: {logits.size()}", flush=True) loss = F.nll_loss(logits[:args.batch_size], data.y[batch_vtxs].long()) # SAGEConv # loss = F.nll_loss(logits[:args.batch_size], data.y[seeds].long()) # quiver @@ -1453,7 +1542,7 @@ def main(args, batches=None): # 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.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): + if args.sample_method == "sage" and args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): model = model.cpu() train_nid = train_nid.cpu() test_nid = test_nid.cpu() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index f84f81f..4e6e0b0 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -1335,6 +1335,7 @@ __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; @@ -1343,7 +1344,7 @@ __device__ int binary_searchd(double *arr, double val, int imin, int imax) { imax = imid - 1; } } - + return ans; } @@ -2237,6 +2238,7 @@ __global__ void ThrowDarts1D(double *dart_values, double *ps_p_values, int *h_va 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); // } From deda84203952142ccb040ecf99952f8c29517d51 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 31 Oct 2023 14:34:18 -0700 Subject: [PATCH 168/178] Last commit of MLSys code --- cagnet/samplers/ladies.py | 1 + cagnet/samplers/sage.py | 40 +- cagnet/samplers/utils.py | 374 ++-- examples/gcn_15d.py | 46 +- sparse-extension/sparse_coo_tensor.cu | 2706 +++++++++++++------------ 5 files changed, 1550 insertions(+), 1617 deletions(-) diff --git a/cagnet/samplers/ladies.py b/cagnet/samplers/ladies.py index 60e94ff..9cb252e 100644 --- a/cagnet/samplers/ladies.py +++ b/cagnet/samplers/ladies.py @@ -52,6 +52,7 @@ def ladies_sampler(adj_matrix, batches, batch_size, frontier_size, mb_count_tota 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: diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 892c917..5d628b5 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -58,6 +58,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total 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) @@ -83,15 +84,14 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_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 - if baseline_compare: - total_start_timer.record() - for i in range(n_layers): - print(f"Sampling layer {i}", flush=True) + start_time(total_start_timer) + # print(f"Sampling layer {i}", flush=True) if i == 0: nnz = batch_size else: @@ -99,6 +99,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # 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)) @@ -116,10 +117,9 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total nnz = batch_size * int(np.prod(frontier_sizes[:i], dtype=int)) if not replicate_graph: current_frontier = current_frontier.to_sparse_csr() + timing_dict["sage-startiter"].append(stop_time(start_timer, stop_timer)) # Expand batches matrix - if baseline_compare: - total_start_timer.record() 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", @@ -128,8 +128,10 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # dist.barrier() # adj_matrix = adj_matrix.to_sparse_coo() - if p.layout == torch.sparse_csr: - p = p.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] @@ -204,27 +206,25 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # frontiers[i + 1] = next_frontier_select.clone() if i + 1 == n_layers: frontiers[i + 1] = next_frontier_select - else: - del next_frontier_select - del p + # else: + # del next_frontier_select + # del p - if i > 0: - del current_frontier_select + # 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)) - # print(f"total_time: {stop_time(total_start_timer, total_stop_timer)}", flush=True) - if baseline_compare: - total_stop_timer.record() - torch.cuda.synchronize() - total_time = total_start_timer.elapsed_time(total_stop_timer) - print(f"total_time: {total_time}", 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" or k == "frontier-row-col-select" or k == "adj-row-col-select" or k.startswith("sample"): + 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): diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index c2c7ad9..c9ee900 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -7,11 +7,11 @@ 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, \ + 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 + rearrangel_rows_gpu, reduce_sum_gpu, sum_csrd_gpu, sum_csri_gpu timing = True @@ -43,7 +43,7 @@ def stop_time_add(start_timer, stop_timer, timing_dict, range_name, barrier=Fals 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, timing_dict=None): +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) @@ -89,28 +89,35 @@ def csr_allreduce(mat, left, right, rank, name=None, timing_dict=None): # 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) - # mat = mat + mat_recv - - # 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 = torch.sparse_csr_tensor(mat_sum_crows, mat_sum_cols, mat_sum_vals, mat.size()) + + 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: @@ -600,6 +607,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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] = \ @@ -684,6 +692,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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: @@ -694,6 +703,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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() @@ -782,7 +792,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups # 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": + 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)) @@ -809,7 +819,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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] + 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 @@ -817,7 +827,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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, matb_recv_cols, matc_chunk_cols) + 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) @@ -870,9 +880,32 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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) @@ -922,7 +955,7 @@ def dist_saspgemm15D(mata, matb, replication, rank, size, row_groups, col_groups 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", timing_dict) + 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) @@ -1058,14 +1091,17 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication adj_matrix.crow_indices(), 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_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_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) @@ -1087,22 +1123,25 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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) - # scatterd_add_gpu(p_den, p_num_indices[0, :], p_num_values, p_num_values.size(0)) - # frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) - p_den.scatter_add_(0, p_num_indices[0, :].long(), p_num_values) + # 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()) - 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])) + torch.cuda.synchronize() timing_dict["compute-p"].append(stop_time(start_timer, stop_timer)) return p @@ -1133,7 +1172,14 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, # n_darts_col = int(avg_degree * frontier_size / avg_degree) n_darts_col = frontier_size - next_frontier = torch.sparse_coo_tensor(indices=p._indices(), + 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), @@ -1143,12 +1189,17 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, frontier_nnz_sizes = torch.cuda.IntTensor(p.size(0)).fill_(0) ones = torch.cuda.IntTensor(next_frontier._indices()[0,:].size(0)).fill_(1) - frontier_nnz_sizes.scatter_add_(0, next_frontier._indices()[0,:], ones) + 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 - zero_idxs = p._indices()[0, (p._values() == 0)].unique() - frontier_nnz_sizes[zero_idxs] = 0 + # 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) @@ -1177,9 +1228,10 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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) + ps_p_values = torch.cumsum(p.values(), dim=0) p_rowsum.fill_(0) - p_rowsum.scatter_add_(0, p._indices()[0, :], p._values()) + # 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 @@ -1244,11 +1296,14 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, sampled_count_old = sampled_count.clone() sampled_count = sampled_count.fill_(0) sampled_count = torch.clamp(frontier_size - frontier_nnz_sizes, min=0) - 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) + 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 @@ -1371,216 +1426,24 @@ def sample(p, frontier_size, mb_count, node_count_total, n_darts, replication, 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,:]) - p._values()[filled_minibatches_mask] = 0 - timing_dict["sample-set-probs"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() - next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] - 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_nnzidxs, \ - next_frontier._values()[next_frontier_nnzvals].int()) - 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])) - 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 - timing_dict["sample-gather-counts-preproc"].append(stop_time(start_timer, stop_timer)) - - start_time(start_timer) - # dist.all_gather_into_tensor(dart_hits_count_recv, dart_hits_count_select, group=row_groups[rank_c]) - # dist.all_gather_into_tensor(sampled_count_recv, sampled_count, group=row_groups[rank_c]) - # dart_hits_count = dart_hits_count_recv[(dart_hits_count_recv != -1)] - # sampled_count = sampled_count_recv[(sampled_count_recv != -1)] - dist.all_gather_into_tensor(dart_hits_count, dart_hits_count_select, group=row_groups[rank_c]) - dist.all_gather_into_tensor(sampled_count, sampled_count, 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 + p.values()[dart_hits_mask] = 0.0 filled_minibatches = sampled_count == frontier_size - filled_minibatches_mask = torch.gather(filled_minibatches, 0, p._indices()[0,:]) - p._values()[filled_minibatches_mask] = 0 + # 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) - next_frontier_nnzvals = next_frontier._values().nonzero().squeeze() - next_frontier_nnzidxs = next_frontier._indices()[0,next_frontier_nnzvals] 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_nnzidxs, \ - next_frontier._values()[next_frontier_nnzvals].int()) + 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)) @@ -1588,9 +1451,6 @@ def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ timing_dict["sampling-iters"].append(stop_time(sample_start_timer, sample_stop_timer)) - # print(f"iter_count: {iter_count}") - # print(f"selection_iter_count: {selection_iter_count}") - return next_frontier def select(next_frontier, adj_matrix, batches, sa_masks, nnz, \ diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 49e6607..2dd46cb 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -793,19 +793,27 @@ def main(args, batches=None): # dist.send(test_idx_len, dst=dst_gpu) dist.send(train_idx, dst=dst_gpu) # dist.send(test_idx, dst=dst_gpu) - if args.replicate_graph: - edge_count = g_loc.col_indices().size(0) + 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: - 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(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()) - print("host2device", flush=True) - g_loc = g_loc.to(device) features_loc = features_loc.to(device) # features_loc, g_loc, _, _ = one5d_partition(rank, size, inputs, adj_matrix, data, inputs, num_classes, \ @@ -915,12 +923,19 @@ def main(args, batches=None): g_loc = torch.sparse_coo_tensor(g_loc_indices, g_loc_values, \ torch.Size([g_loc_meta[1], g_loc_meta[2]])) - 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) + 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": @@ -1161,9 +1176,9 @@ def main(args, batches=None): # nnz_recv_upperbound = edge_count // (size // args.replication) + print("Sampling", flush=True) if epoch >= 1: start_time(start_timer, timing_arg=True) - print("Sampling", flush=True) torch.cuda.nvtx.range_push("nvtx-sampling") if args.sample_method == "ladies": current_frontier, next_frontier, adj_matrices_bulk = \ @@ -1338,7 +1353,6 @@ def main(args, batches=None): if epoch >= 1: start_time(start_inner_timer, timing_arg=True) torch.cuda.nvtx.range_push("nvtx-selectfeats") - print(f"adjs: {adjs}", flush=True) if size > 1: src_vtxs_sort = torch.cuda.LongTensor(src_vtxs.size(0)) @@ -1542,7 +1556,7 @@ def main(args, batches=None): # 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 False and args.sample_method == "sage" and args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): model = model.cpu() train_nid = train_nid.cpu() test_nid = test_nid.cpu() diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 4e6e0b0..7bca861 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -62,1274 +62,1257 @@ using namespace at::sparse; #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); -// } -// -// 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); -// } -// -// 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); -// } -// -// __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) { -// // printf("infinite loop key %d adr %d offset: %d hash: %d d_check[adr] %d bcol: %d max_row_nz: %d\n", key, adr, offset, hash, d_check[adr], bcol, max_row_nz); -// 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); -// } -// -// 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])); -// // torch::Tensor d_fail_count_tens = torch::zeros({1}, torch::TensorOptions() -// // .dtype(torch::kInt32) -// // .device(torch::kCUDA)); -// // torch::Tensor d_fail_perm_tens = torch::zeros({bin->bin_size[i]}, torch::TensorOptions() -// // .dtype(torch::kInt32) -// // .device(torch::kCUDA)); -// // int *d_fail_count = d_fail_count_tens.data(); -// // int *d_fail_perm = d_fail_perm_tens.data(); -// 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)); -// -// // torch::Tensor d_check_tens = torch::zeros({(long)table_size}, torch::TensorOptions() -// // .dtype(torch::kInt32) -// // .device(torch::kCUDA)); -// // int *d_check = d_check_tens.data(); -// 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)); -// -// torch::Tensor d_check_tens = torch::zeros({table_size}, torch::TensorOptions() -// .dtype(torch::kInt32) -// .device(torch::kCUDA)); -// torch::Tensor d_value_tens = torch::zeros({table_size}, torch::TensorOptions() -// .dtype(torch::kFloat32) -// .device(torch::kCUDA)); -// int *d_check = d_check_tens.data(); -// float *d_value = d_value_tens.data(); // assumes real type is float -// -// 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(); -// } +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; @@ -1538,12 +1521,8 @@ void spmm_gpu(const at::Tensor& A_rowindices, clock_t start, stop; - // int32_t *d_a_csrrows; - // cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); - torch::Tensor A_csrrows_tens = torch::zeros({n + 1}, torch::TensorOptions() - .dtype(torch::kInt32) - .device(torch::kCUDA)); - int32_t *d_a_csrrows = A_csrrows_tens.data(); + int32_t *d_a_csrrows; + cudaMalloc(&d_a_csrrows, (n + 1) * sizeof(int32_t)); CHECK_CUSPARSE(cusparseXcoo2csr(handle, A_rowindices.data(), @@ -1623,7 +1602,7 @@ void spmm_gpu(const at::Tensor& A_rowindices, d_buffer)); // buffer - // cudaFree(d_a_csrrows); + cudaFree(d_a_csrrows); cudaFree(d_buffer); // Column-major to row-major @@ -1802,68 +1781,74 @@ std::vector nsparse_spgemm( int32_t k, int32_t m) { - // sfBIN bin; + 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(); + /* 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(); + 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); + 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(); + // 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); + 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); - - // /* Set bin */ - // set_min_bin(&bin, n); + 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)); + // 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(); + 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); + /* 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); + release_bin(bin); + std::cout << "after release_bin" << std::endl; - // CHECK_ERROR("nsparse_spgemm") + CHECK_ERROR("nsparse_spgemm") - // return {C_crow_indices, C_col_indices, C_values}; - // // return {c_csr_offsets, c_columns, c_values}; + return {C_crow_indices, C_col_indices, C_values}; + // return {c_csr_offsets, c_columns, c_values}; } std::vector spgemm_gpu( @@ -2341,6 +2326,76 @@ void normalize_gpu(const at::Tensor& output, const at::Tensor& input, const at:: 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; @@ -2834,6 +2889,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); From 66d22d8a8b9d1b474fc47dc57a2d29de1f6c4099 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 10 Nov 2023 19:21:38 -0800 Subject: [PATCH 169/178] Linear scaling optimizer adjustment --- examples/gcn_15d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 2dd46cb..6585a48 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -1066,7 +1066,7 @@ def main(args, batches=None): model = model.to(device) # use optimizer - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) dur = [] @@ -1556,7 +1556,7 @@ def main(args, batches=None): # print(f"adj1_nnz.median: {np.median(adj1_nnzs)}") # print(f"adj_nnzs.median: {np.median(adj_nnzs)}") print(f"total: {np.sum(dur) / epoch}", flush=True) - if False and args.sample_method == "sage" and args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): + if args.sample_method == "sage" and args.dataset != "ogbn-papers100M" and args.dataset != "Amazon" and ("Protein" not in args.dataset): model = model.cpu() train_nid = train_nid.cpu() test_nid = test_nid.cpu() From a95d58af8085f2f1f3e11b4bcface8b12cc720f5 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 4 Dec 2023 15:12:38 -0800 Subject: [PATCH 170/178] Adding wandb --- .gitignore | 1 + examples/gcn_15d.py | 17 +++++++++++++---- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index 84de964..666d777 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,4 @@ examples/*.batch *.out *.nsys-rep *.ncu-rep +examples/wandb diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 6585a48..5bf7995 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -29,6 +29,8 @@ 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, @@ -75,10 +77,11 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) # # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) - self.layers.append(SAGEConv(in_feats, n_hidden, root_weight=False, bias=False)) + self.layers.append(SAGEConv(in_feats, n_hidden)) for _ in range(n_layers - 2): - self.layers.append(SAGEConv(n_hidden, n_hidden, root_weight=False, bias=False)) - self.layers.append(SAGEConv(n_hidden, n_classes, root_weight=False, bias=False)) + self.layers.append(SAGEConv(n_hidden, n_hidden)) + self.layers.append(SAGEConv(n_hidden, n_classes)) + # self.layers.append(SAGEConv(in_feats, n_classes, root_weight=False, bias=False)) def forward(self, graphs, inputs, epoch): h = inputs @@ -1065,6 +1068,9 @@ def main(args, batches=None): col_groups=col_groups) model = model.to(device) + for name, W in model.named_parameters(): + print(f"name: {name} W.sum: {W.sum()}", flush=True) + # use optimizer optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) @@ -1461,6 +1467,7 @@ def main(args, batches=None): # adjs = [adj.to(rank) for adj in adjs] # features_batch = features_loc[n_id] optimizer.zero_grad() + print(f"adjs: {adjs} features_batch: {features_batch}", flush=True) logits = model(adjs, features_batch, epoch) torch.cuda.nvtx.range_pop() # nvtx-fwd if epoch >= 1: @@ -1470,7 +1477,8 @@ def main(args, batches=None): 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].long()) # SAGEConv + loss = F.nll_loss(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 @@ -1570,6 +1578,7 @@ def main(args, batches=None): test_nid = test_nid.to(device) model = model.to(device) print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) + # 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) From cf822fb4c96ad47c0b9ad6178fd44d29ca770047 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Mon, 22 Jan 2024 01:04:20 -0800 Subject: [PATCH 171/178] Cora accuracy fixes --- cagnet/samplers/sage.py | 87 +++++++- examples/cora_nl1_rw_bias_correct_torchviz | 37 ++++ .../cora_nl1_rw_bias_correct_torchviz.svg | 189 ++++++++++++++++++ examples/cora_nl1_rw_bias_torchviz | 37 ++++ examples/cora_nl1_rw_bias_torchviz.svg | 189 ++++++++++++++++++ examples/gcn_15d.py | 179 +++++++++++++++-- examples/sage_conv.py | 166 +++++++++++++++ tests/test_sage.py | 60 ++++-- 8 files changed, 896 insertions(+), 48 deletions(-) create mode 100644 examples/cora_nl1_rw_bias_correct_torchviz create mode 100644 examples/cora_nl1_rw_bias_correct_torchviz.svg create mode 100644 examples/cora_nl1_rw_bias_torchviz create mode 100644 examples/cora_nl1_rw_bias_torchviz.svg create mode 100644 examples/sage_conv.py diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 5d628b5..4b0ce87 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -93,7 +93,11 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total start_time(total_start_timer) # print(f"Sampling layer {i}", flush=True) if i == 0: - nnz = batch_size + 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() @@ -114,7 +118,16 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total 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 = 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)) @@ -138,17 +151,26 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total 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() - frontier_nnz_sizes = torch.histc(next_frontier._indices()[0,next_frontier_nnz], bins=p.size(0)) + # 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 * 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) @@ -167,9 +189,12 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # size=(nnz * mb_count, node_count_total)) # # size=(batch_size * mb_count, node_count_total)) next_frontier = sparse_coo_tensor_gpu(next_frontier_idxs, next_frontier_values, - torch.Size([nnz * mb_count, node_count_total])) + torch.Size([nnz_row, node_count_total])) + # torch.Size([nnz * mb_count, node_count_total])) - next_frontier_select = next_frontier._indices()[1,:].view(mb_count * nnz, frontier_size) + 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)) @@ -186,13 +211,13 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total batch_vals = torch.cuda.LongTensor(current_frontier_select.size()).fill_(1) next_frontier_select_vals = next_frontier._values().view(-1) - batch_rows = torch.arange(mb_count * nnz, device=gpu).view(mb_count * nnz, 1) next_frontier_select_rows = next_frontier._indices()[0,:].view(-1) nnz_mask = next_frontier_select_vals.nonzero().squeeze() adj_sample_rows = next_frontier_select_rows[nnz_mask] adj_sample_cols = torch.arange(next_frontier_select.numel(), device=gpu) - adj_sample_cols = adj_sample_cols.remainder(next_frontier_select.size(1) * nnz) + # adj_sample_cols = adj_sample_cols.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) @@ -201,11 +226,53 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # adj_matrix_sample = torch.sparse_coo_tensor(adj_matrices_indices, adj_matrices_values, # size=torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) adj_matrix_sample = sparse_coo_tensor_gpu(adj_matrices_indices, adj_matrices_values, - torch.Size([mb_count * nnz, next_frontier_select.size(1) * nnz])) + # 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) + print(f"nnz_row: {nnz_row} nnz_col: {nnz_col}", flush=True) + print(f"next_frontier_indices.size: {next_frontier._indices().size()} next_frontier_increments.size: {next_frontier_increments.size()}", flush=True) + 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"current_frontier: {current_frontier}", flush=True) + print(f"before next_frontier: {next_frontier}", flush=True) + next_frontier = current_frontier + next_frontier + print(f"after1 next_frontier: {next_frontier}", flush=True) + next_frontier = next_frontier.coalesce() + print(f"after2 next_frontier: {next_frontier}", flush=True) + + # next_frontier_missing_vtxs = next_frontier._values() == 0 + next_frontier_select = next_frontier._indices()[1,:].view(-1) #.clone() + # 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) + print(f"adj_sample: {adj_matrix_sample}", flush=True) if i + 1 == n_layers: - frontiers[i + 1] = next_frontier_select + frontiers[i + 1] = next_frontier_select.view(-1, 1) # else: # del next_frontier_select # del p 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/gcn_15d.py b/examples/gcn_15d.py index 5bf7995..96ccb89 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -13,7 +13,8 @@ import torch_geometric import torch_geometric.transforms as T from torch_geometric.data import Data, Dataset -from torch_geometric.nn import SAGEConv +# from torch_geometric.nn import SAGEConv +from sage_conv import SAGEConv from torch_geometric.datasets import Planetoid, Reddit from torch_geometric.loader import NeighborSampler from torch_geometric.utils import add_remaining_self_loops @@ -81,7 +82,7 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa for _ in range(n_layers - 2): self.layers.append(SAGEConv(n_hidden, n_hidden)) self.layers.append(SAGEConv(n_hidden, n_classes)) - # self.layers.append(SAGEConv(in_feats, n_classes, root_weight=False, bias=False)) + # self.layers.append(SAGEConv(in_feats, n_classes)) def forward(self, graphs, inputs, epoch): h = inputs @@ -91,14 +92,17 @@ def forward(self, graphs, inputs, epoch): # 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) + print(f"h.sum: {h.sum()}", flush=True) return h @torch.no_grad() @@ -127,10 +131,10 @@ def evaluate(self, graph, features, test_idx, labels): # return features subgraph_loader = NeighborSampler(graph, node_idx=None, - sizes=[-1], batch_size=2048, + sizes=[-1], batch_size=1024, # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, # sizes=[-1, -1], batch_size=512, - shuffle=False) + shuffle=False, num_workers=6) non_eval_timings = copy.deepcopy(self.timings) for l, layer in enumerate(self.layers): hs = [] @@ -675,6 +679,12 @@ def main(args, batches=None): # 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)) @@ -1035,7 +1045,6 @@ def main(args, batches=None): # size=(node_count, node_count)) # adj_matrix = adj_matrix.coalesce().cuda().double() - # return frontiers, adj_matrices, adj_matrix # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model @@ -1072,7 +1081,8 @@ def main(args, batches=None): print(f"name: {name} W.sum: {W.sum()}", flush=True) # use optimizer - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) + # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) dur = [] @@ -1096,6 +1106,8 @@ def main(args, batches=None): # g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), \ # g_loc.size()) g_loc = g_loc.to(device) + print(f"g_loc: {g_loc}", flush=True) + print(f"g_loc[vtx23]: {g_loc.col_indices()[g_loc.crow_indices()[23]:g_loc.crow_indices()[24]]}", flush=True) # csr_topo = quiver.CSRTopo(adj_matrix) # quiver_sampler = quiver.pyg.GraphSageSampler(csr_topo, [args.samp_num] * args.n_layers, 0, # mode='UVA') @@ -1119,7 +1131,7 @@ def main(args, batches=None): 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) + # 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 @@ -1146,6 +1158,7 @@ def main(args, batches=None): args.batch_size = batch_count - b torch.cuda.nvtx.range_push("nvtx-partition-batches") batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) + print(f"batches: {batches}", flush=True) if args.sample_method == "sage": batches_loc = one5d_partition_mb(rank, size, batches, 1, args.n_bulkmb) elif args.sample_method == "ladies": @@ -1211,6 +1224,7 @@ def main(args, batches=None): col_groups, args.timing, args.baseline, args.replicate_graph) + print(f"frontiers_bulk: {frontiers_bulk}", flush=True) if epoch >= 1: model.timings["sample"].append(stop_time(start_timer, stop_timer, timing_arg=True)) torch.cuda.nvtx.range_pop() # nvtx-sampling @@ -1222,6 +1236,12 @@ def main(args, batches=None): 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) @@ -1234,9 +1254,26 @@ def main(args, batches=None): 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[:-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 + + print(f"i: {i} src_select_min: {src_select_min} src_select_max: {src_select_max}", flush=True) src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) + print(f"i: {i} batch_vtxs: {batch_vtxs}", flush=True) + print(f"i: {i} src_vtxs: {src_vtxs}", flush=True) + + # 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] @@ -1246,9 +1283,22 @@ def main(args, batches=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_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] + + print(f"i: {i} nnz_row: {nnz_row} nnz_col: {nnz_col}", flush=True) + 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)) + print(f"min: {adj_row_select_min} max: {adj_row_select_max}", flush=True) if epoch >= 1: start_time(start_inner_timer, timing_arg=True) @@ -1256,6 +1306,7 @@ def main(args, batches=None): 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] + print(f"adj_mat: {adj_matrices_bulk[l]}", flush=True) 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() @@ -1302,8 +1353,12 @@ def main(args, batches=None): # args.batch_size * (np.prod(args.samp_num[:(l+1)], dtype=int)))) # adj_nnzs.append(adj_sample_values.size(0)) if l > 0: + print(f"adj_sample_crows: {adj_sample_crows}", flush=True) adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] # adj_sample_row_lens[adj_sample_skip_cols] = 0 + print(f"row_lens: {adj_sample_row_lens}", flush=True) + print(f"skip_cols: {adj_sample_skip_cols}", flush=True) + print(F"adj_sample_cols: {adj_sample_cols}", flush=True) 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] @@ -1327,21 +1382,41 @@ def main(args, batches=None): # 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 = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) + col_count = nnz_row * args.samp_num[l] + print(f"before adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + 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) + print(f"adj_sample_cols: {adj_sample_cols} skip_cols: {adj_sample_skip_cols} col_count: {col_count}", flush=True) + print(f"skip_cols_prev: {adj_sample_skip_cols_prev}", flush=True) + 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 + print(f"inc: {inc}", flush=True) + 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"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) + # col_count = args.batch_size * np.prod(args.samp_num[:(l+1)], dtype=int) - \ + # adj_sample_skip_cols.size(0) + 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)) + print(f"row_count: {row_count} col_count: {col_count} nnz_col: {nnz_col} skip_cols: {adj_sample_skip_cols.size(0)}", flush=True) adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ - adj_sample_values, (row_count, col_count)) + 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)) @@ -1361,6 +1436,10 @@ def main(args, batches=None): torch.cuda.nvtx.range_push("nvtx-selectfeats") if size > 1: + print(f"batch_vtxs: {batch_vtxs}", flush=True) + print(f"src_vtxs: {src_vtxs}", flush=True) + src_vtxs = torch.cat((batch_vtxs, src_vtxs)) + print(f"src_vtxs2: {src_vtxs}", flush=True) 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] @@ -1403,19 +1482,68 @@ def main(args, batches=None): input_tally, output_tally, col_groups[rank_col]) - #features_batch = torch.cuda.FloatTensor(output_features.size()) + # 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_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)) else: + print(f"batch_vtxs.size: {batch_vtxs.size()}", flush=True) + print(f"src_vtxs.size: {src_vtxs.size()}", flush=True) + print(f"batch_vtxs: {batch_vtxs}", flush=True) + print(f"src_vtxs: {src_vtxs}", flush=True) + # adj_sample_skip_cols += args.batch_size # manual fix, change this + print(f"(src_vtxs == 0).nnz: {(src_vtxs == 0).nonzero().squeeze()}", flush=True) + print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) + + # 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) - features_mask[adj_sample_skip_cols] = False + # 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() + print(f"src_vtxs == 0: {(src_vtxs == 0).nonzero().squeeze()}", flush=True) + print(f"adj_sample_skip_cols_mod: {adj_sample_skip_cols}", flush=True) + print(f"adj_sample_skip_cols_feats: {adj_sample_skip_cols_feats}", flush=True) + print(f"features_batch.size: {features_batch.size()}", flush=True) + features_mask[adj_sample_skip_cols_feats] = False features_batch = features_batch[features_mask] torch.cuda.nvtx.range_pop() # nvtx-selectfeats @@ -1468,7 +1596,10 @@ def main(args, batches=None): # features_batch = features_loc[n_id] optimizer.zero_grad() print(f"adjs: {adjs} features_batch: {features_batch}", flush=True) + print(f"features_batch.size: {features_batch.size()}", flush=True) logits = model(adjs, features_batch, epoch) + print(f"logits.size: {logits.size()}", flush=True) + print(f"logits.sum: {logits.sum()}", flush=True) 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)) @@ -1493,16 +1624,19 @@ def main(args, batches=None): if epoch >= 1: model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) + for name, W in model.named_parameters(): + print(f"name: {name} W.sum: {W.sum()} W.grad.sum: {W.grad.sum()}", flush=True) for W in model.parameters(): dist.all_reduce(W.grad) W.grad /= size - torch.cuda.nvtx.range_push("nvtx-optstep") optimizer.step() + torch.cuda.nvtx.range_push("nvtx-optstep") - for i in range(len(adjs)): + 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)) @@ -1521,6 +1655,11 @@ def main(args, batches=None): 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) @@ -1574,10 +1713,10 @@ def main(args, batches=None): # acc1 = int(res[train_nid].sum()) / train_nid.size(0) acc3 = int(res[test_nid].sum()) / test_nid.size(0) # acc3 = correct_count / test_nid.size(0) + print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) train_nid = train_nid.to(device) test_nid = test_nid.to(device) model = model.to(device) - print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {:.4f} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) # 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) 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/tests/test_sage.py b/tests/test_sage.py index 4fb07ea..681688e 100644 --- a/tests/test_sage.py +++ b/tests/test_sage.py @@ -17,6 +17,7 @@ def main(args): # 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") @@ -39,6 +40,8 @@ def main(args): 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] @@ -108,39 +111,60 @@ def main(args): 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=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("--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('--cuda', type=int, default=0, - help='Avaiable GPU ID') - parser.add_argument('--sample-method', type=str, default='sage', - help='Sampled Algorithms: ladies/fastgcn/full') - parser.add_argument('--batch-num', type=int, default= 10, - help='Maximum Batch Number') - parser.add_argument('--pool-num', type=int, default= 10, - help='Number of Pool') - parser.add_argument('--n-bulkmb', type=int, default=1, - help='Number of minibatches to sample in bulk') - parser.add_argument('--n-darts', type=int, default=-1, - help='Number of darts to throw per minibatch in LADIES sampling') + parser.add_argument('--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) From 03e36634aecfb3397b99d052368272467a1dd5e2 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 30 Mar 2024 00:22:24 -0700 Subject: [PATCH 172/178] Added frontier adding kernel --- cagnet/samplers/sage.py | 41 ++++++++---- examples/gcn_15d.py | 47 +------------- sparse-extension/sparse_coo_tensor.cu | 94 +++++++++++++++++++++++++++ 3 files changed, 126 insertions(+), 56 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 4b0ce87..52ca9ce 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -11,7 +11,7 @@ compute_darts1d_gpu, throw_darts1d_gpu, normalize_gpu, \ shift_rowselect_gpu, shift_colselect_gpu, \ scatteri_add_gpu, rowselect_coo_gpu, \ - sparse_coo_tensor_gpu + sparse_coo_tensor_gpu, frontier_sum_gpu timing = True @@ -192,7 +192,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total torch.Size([nnz_row, node_count_total])) # torch.Size([nnz * mb_count, node_count_total])) - next_frontier = next_frontier.coalesce() + # 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(), \ @@ -243,23 +243,41 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total 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) - print(f"nnz_row: {nnz_row} nnz_col: {nnz_col}", flush=True) - print(f"next_frontier_indices.size: {next_frontier._indices().size()} next_frontier_increments.size: {next_frontier_increments.size()}", flush=True) 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"current_frontier: {current_frontier}", flush=True) - print(f"before next_frontier: {next_frontier}", flush=True) - next_frontier = current_frontier + next_frontier - print(f"after1 next_frontier: {next_frontier}", flush=True) - next_frontier = next_frontier.coalesce() - print(f"after2 next_frontier: {next_frontier}", flush=True) + # 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()).fill_(123123) + output_vals = torch.cuda.FloatTensor(current_frontier._nnz() + next_frontier._nnz()).fill_(123123.0) + 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._indices()[1,:].view(-1) #.clone() # next_frontier_select[next_frontier_missing_vtxs] = -1 # adj_matrices[i] = adj_matrix_sample @@ -270,7 +288,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # 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) - print(f"adj_sample: {adj_matrix_sample}", flush=True) if i + 1 == n_layers: frontiers[i + 1] = next_frontier_select.view(-1, 1) # else: diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 96ccb89..3b0c6c1 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -13,8 +13,8 @@ import torch_geometric import torch_geometric.transforms as T from torch_geometric.data import Data, Dataset -# from torch_geometric.nn import SAGEConv -from sage_conv import SAGEConv +from torch_geometric.nn import SAGEConv +# from sage_conv import SAGEConv from torch_geometric.datasets import Planetoid, Reddit from torch_geometric.loader import NeighborSampler from torch_geometric.utils import add_remaining_self_loops @@ -102,7 +102,6 @@ def forward(self, graphs, inputs, epoch): # h = CAGF.log_softmax(self, h, self.partitioning, dim=1) h = F.log_softmax(h, dim=1) - print(f"h.sum: {h.sum()}", flush=True) return h @torch.no_grad() @@ -1048,7 +1047,7 @@ def main(args, batches=None): # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model - # torch.manual_seed(0) + torch.manual_seed(0) if args.sample_method == "sage": model = GCN(num_features, args.n_hidden, @@ -1077,9 +1076,6 @@ def main(args, batches=None): col_groups=col_groups) model = model.to(device) - for name, W in model.named_parameters(): - print(f"name: {name} W.sum: {W.sum()}", flush=True) - # use optimizer # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) @@ -1106,8 +1102,6 @@ def main(args, batches=None): # g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc.values(), \ # g_loc.size()) g_loc = g_loc.to(device) - print(f"g_loc: {g_loc}", flush=True) - print(f"g_loc[vtx23]: {g_loc.col_indices()[g_loc.crow_indices()[23]:g_loc.crow_indices()[24]]}", flush=True) # csr_topo = quiver.CSRTopo(adj_matrix) # quiver_sampler = quiver.pyg.GraphSageSampler(csr_topo, [args.samp_num] * args.n_layers, 0, # mode='UVA') @@ -1158,7 +1152,6 @@ def main(args, batches=None): args.batch_size = batch_count - b torch.cuda.nvtx.range_push("nvtx-partition-batches") batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) - print(f"batches: {batches}", flush=True) if args.sample_method == "sage": batches_loc = one5d_partition_mb(rank, size, batches, 1, args.n_bulkmb) elif args.sample_method == "ladies": @@ -1224,7 +1217,6 @@ def main(args, batches=None): col_groups, args.timing, args.baseline, args.replicate_graph) - print(f"frontiers_bulk: {frontiers_bulk}", flush=True) if epoch >= 1: model.timings["sample"].append(stop_time(start_timer, stop_timer, timing_arg=True)) torch.cuda.nvtx.range_pop() # nvtx-sampling @@ -1265,10 +1257,7 @@ def main(args, batches=None): src_select_min = i * nnz_col src_select_max = (i + 1) * nnz_col - print(f"i: {i} src_select_min: {src_select_min} src_select_max: {src_select_max}", flush=True) src_vtxs = frontiers_bulk[-1][src_select_min:src_select_max,:].view(-1) - print(f"i: {i} batch_vtxs: {batch_vtxs}", flush=True) - print(f"i: {i} src_vtxs: {src_vtxs}", flush=True) # print(f"i: {i} batch_vtxs: {batch_vtxs}", flush=True) # frontiers_test[0][i] = frontiers_bulk[0][batch_select_min:batch_select_max,:] @@ -1306,7 +1295,6 @@ def main(args, batches=None): 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] - print(f"adj_mat: {adj_matrices_bulk[l]}", flush=True) 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() @@ -1353,12 +1341,8 @@ def main(args, batches=None): # args.batch_size * (np.prod(args.samp_num[:(l+1)], dtype=int)))) # adj_nnzs.append(adj_sample_values.size(0)) if l > 0: - print(f"adj_sample_crows: {adj_sample_crows}", flush=True) adj_sample_row_lens = adj_sample_crows[1:] - adj_sample_crows[:-1] # adj_sample_row_lens[adj_sample_skip_cols] = 0 - print(f"row_lens: {adj_sample_row_lens}", flush=True) - print(f"skip_cols: {adj_sample_skip_cols}", flush=True) - print(F"adj_sample_cols: {adj_sample_cols}", flush=True) 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] @@ -1384,12 +1368,9 @@ def main(args, batches=None): # 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] - print(f"before adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) 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) - print(f"adj_sample_cols: {adj_sample_cols} skip_cols: {adj_sample_skip_cols} col_count: {col_count}", flush=True) - print(f"skip_cols_prev: {adj_sample_skip_cols_prev}", flush=True) if l == 0: adj_sample_skip_cols += args.batch_size adj_sample_skip_cols_feats = adj_sample_skip_cols @@ -1398,10 +1379,8 @@ def main(args, batches=None): for adj in adjs: if adj is not None: inc += adj._indices()[1,-1] + 1 - print(f"inc: {inc}", flush=True) 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"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) @@ -1410,7 +1389,6 @@ def main(args, batches=None): # 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)) - print(f"row_count: {row_count} col_count: {col_count} nnz_col: {nnz_col} skip_cols: {adj_sample_skip_cols.size(0)}", flush=True) adj_matrix_sample = torch.sparse_coo_tensor(adj_sample_indices, \ adj_sample_values, (nnz_row, nnz_col)) # (args.batch_size * (args.samp_num ** l), \ @@ -1436,10 +1414,7 @@ def main(args, batches=None): torch.cuda.nvtx.range_push("nvtx-selectfeats") if size > 1: - print(f"batch_vtxs: {batch_vtxs}", flush=True) - print(f"src_vtxs: {src_vtxs}", flush=True) src_vtxs = torch.cat((batch_vtxs, src_vtxs)) - print(f"src_vtxs2: {src_vtxs}", flush=True) 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] @@ -1501,13 +1476,7 @@ def main(args, batches=None): adjs[j] = torch.sparse_coo_tensor(adjs_indices, adjs[j]._values(), size=(adjs[j].size(0), adjs[j].size(1) + args.batch_size)) else: - print(f"batch_vtxs.size: {batch_vtxs.size()}", flush=True) - print(f"src_vtxs.size: {src_vtxs.size()}", flush=True) - print(f"batch_vtxs: {batch_vtxs}", flush=True) - print(f"src_vtxs: {src_vtxs}", flush=True) # adj_sample_skip_cols += args.batch_size # manual fix, change this - print(f"(src_vtxs == 0).nnz: {(src_vtxs == 0).nonzero().squeeze()}", flush=True) - print(f"adj_sample_skip_cols: {adj_sample_skip_cols}", flush=True) # features_batch_vtxs = features_loc[batch_vtxs] # features_batch = torch.cat((features_batch_vtxs, features_batch)) @@ -1539,10 +1508,6 @@ def main(args, batches=None): # 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() - print(f"src_vtxs == 0: {(src_vtxs == 0).nonzero().squeeze()}", flush=True) - print(f"adj_sample_skip_cols_mod: {adj_sample_skip_cols}", flush=True) - print(f"adj_sample_skip_cols_feats: {adj_sample_skip_cols_feats}", flush=True) - print(f"features_batch.size: {features_batch.size()}", flush=True) features_mask[adj_sample_skip_cols_feats] = False features_batch = features_batch[features_mask] @@ -1595,11 +1560,7 @@ def main(args, batches=None): # adjs = [adj.to(rank) for adj in adjs] # features_batch = features_loc[n_id] optimizer.zero_grad() - print(f"adjs: {adjs} features_batch: {features_batch}", flush=True) - print(f"features_batch.size: {features_batch.size()}", flush=True) logits = model(adjs, features_batch, epoch) - print(f"logits.size: {logits.size()}", flush=True) - print(f"logits.sum: {logits.sum()}", flush=True) 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)) @@ -1624,8 +1585,6 @@ def main(args, batches=None): if epoch >= 1: model.timings["bwd"].append(stop_time(start_inner_timer, stop_inner_timer, timing_arg=True)) - for name, W in model.named_parameters(): - print(f"name: {name} W.sum: {W.sum()} W.grad.sum: {W.grad.sum()}", flush=True) for W in model.parameters(): dist.all_reduce(W.grad) W.grad /= size diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 7bca861..93f3207 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2876,6 +2876,99 @@ void reduce_sum_gpu(const at::Tensor& matc_crows, const at::Tensor& mata_crows, 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"); @@ -2905,4 +2998,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); } From 6a198b5b5bb33a7e95af354765bcc4b540eaaabc Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sat, 30 Mar 2024 00:50:44 -0700 Subject: [PATCH 173/178] Fixed feature fetching for p>1 --- cagnet/samplers/sage.py | 28 ++++++++++++++-------------- examples/gcn_15d.py | 27 ++++++++++++++++++++------- 2 files changed, 34 insertions(+), 21 deletions(-) diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 52ca9ce..c92fb62 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -259,20 +259,20 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # 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()).fill_(123123) - output_vals = torch.cuda.FloatTensor(current_frontier._nnz() + next_frontier._nnz()).fill_(123123.0) - 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 = current_frontier + next_frontier + # output_indices = torch.cuda.LongTensor(2, current_frontier._nnz() + next_frontier._nnz()).fill_(123123) + # output_vals = torch.cuda.FloatTensor(current_frontier._nnz() + next_frontier._nnz()).fill_(123123.0) + # 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() diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 3b0c6c1..b1ae1fd 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -1047,7 +1047,7 @@ def main(args, batches=None): # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model - torch.manual_seed(0) + # torch.manual_seed(0) if args.sample_method == "sage": model = GCN(num_features, args.n_hidden, @@ -1414,7 +1414,7 @@ def main(args, batches=None): torch.cuda.nvtx.range_push("nvtx-selectfeats") if size > 1: - src_vtxs = torch.cat((batch_vtxs, src_vtxs)) + # 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] @@ -1463,18 +1463,31 @@ def main(args, batches=None): 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 + # 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)): + + # 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,:] - adjs_cols += args.batch_size + 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), adjs[j].size(1) + args.batch_size)) + # 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 From 208d8e4d1603b45f3fcef029102d3656de3c0691 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Sun, 7 Apr 2024 01:48:01 -0700 Subject: [PATCH 174/178] Added new evaluate() code and debugging overflow on undirected papers --- cagnet/samplers/sage.py | 45 +++-- cagnet/samplers/utils.py | 17 ++ examples/dgl_papers.py | 223 ++++++++++++++++++++++ examples/gcn_15d.py | 263 ++++++++++++++++---------- examples/pyg_papers.py | 209 ++++++++++++++++++++ sparse-extension/sparse_coo_tensor.cu | 15 +- 6 files changed, 641 insertions(+), 131 deletions(-) create mode 100644 examples/dgl_papers.py create mode 100644 examples/pyg_papers.py diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index c92fb62..8f86ac8 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -96,7 +96,7 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total 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) + # 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 @@ -127,12 +127,15 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # 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) + # 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 + print(f"l: {i} current_frontier: {current_frontier}", flush=True) + print(f"l: {i} current_frontier.indices.min: {current_frontier._indices().min()}", flush=True) + print(f"l: {i} current_frontier.indices.dtype: {current_frontier._indices().dtype}", flush=True) 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", @@ -259,20 +262,20 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # 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()).fill_(123123) - # output_vals = torch.cuda.FloatTensor(current_frontier._nnz() + next_frontier._nnz()).fill_(123123.0) - # 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 = 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() @@ -290,12 +293,12 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total # 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 + else: + del next_frontier_select + del p - # if i > 0: - # del current_frontier_select + if i > 0: + del current_frontier_select current_frontier = next_frontier timing_dict["adj-row-col-select"].append(stop_time(start_timer, stop_timer)) diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index c9ee900..d28be2a 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1080,17 +1080,34 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication 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_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) + print(f"before crow_indices: {adj_matrix.crow_indices()}", flush=True) + print(f"before crow_indices.min: {adj_matrix.crow_indices().min()}", flush=True) + print(f"before crow_indices.max: {adj_matrix.crow_indices().max()}", flush=True) + print(f"before crow_indices.dtype: {adj_matrix.crow_indices().dtype}", flush=True) + print(f"before col_indices: {adj_matrix.col_indices()}", flush=True) + print(f"before col_indices.min: {adj_matrix.col_indices().min()}", flush=True) + print(f"before col_indices.max: {adj_matrix.col_indices().max()}", flush=True) + print(f"before col_indices.dtype: {adj_matrix.col_indices().dtype}", flush=True) + print(f"before matc_chunk_cols.min: {matc_chunk_cols.min()}", flush=True) rearrange_rows_gpu(numerator._indices()[0,:], numerator._indices()[1,:], matc_chunk_crows, adj_matrix.crow_indices(), adj_matrix.col_indices(), matc_chunk_cols) + print(f"after matc_chunk_cols: {matc_chunk_cols}", flush=True) + print(f"after matc_chunk_cols.min: {matc_chunk_cols.min()}", flush=True) # p_num_rows = torch.repeat_interleave( # torch.arange(matc_chunk_crows.size(0) - 1, # dtype=torch.int32, 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 index b1ae1fd..33fa420 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -98,7 +98,7 @@ def forward(self, graphs, inputs, epoch): 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 = 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) @@ -106,98 +106,73 @@ def forward(self, graphs, inputs, epoch): @torch.no_grad() def evaluate(self, graph, features, test_idx, labels): - # subgraph_loader = NeighborSampler(graph, node_idx=None, - # sizes=[-1], batch_size=2048, - # shuffle=False, num_workers=6) - # for i in range(self.n_layers): - # xs = [] - # for batch_size, n_id, adj in subgraph_loader: - # edge_index, _, size = adj.to(self.device) - # x = features[n_id].to(self.device) - # # edge_index, _, size = adj - # # x = features[n_id] - # # x_target = x[:size[1]] - # # x = self.convs[i]((x, x_target), edge_index) - # x = self.layers[i](x, edge_index) - # if i != self.n_layers - 1: - # x = F.relu(x) - # # xs.append(x) - # xs.append(x[:batch_size]) - - # features = torch.cat(xs, dim=0) - - # return features - - subgraph_loader = NeighborSampler(graph, node_idx=None, - sizes=[-1], batch_size=1024, - # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, - # sizes=[-1, -1], batch_size=512, - shuffle=False, num_workers=6) - 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 + # for papers100M + subgraph_loader = NeighborSampler(graph, 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): + edge_index, _, size = adjs[l] + # edge_index_cpu, _, size = adjs[l] + # del edge_index_cpu adj_batch = torch.sparse_coo_tensor(edge_index, torch.FloatTensor(edge_index.size(1)).fill_(1.0), + # torch.cuda.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) + + del edge_index + del adj_batch 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) + hs.append(h[:batch_size]) + batch_count += 1 + features = torch.cat(hs, dim=0) return features - - # print(f"test_idx: {test_idx}", flush=True) - # non_eval_timings = copy.deepcopy(self.timings) - # correct_count = 0 - # for batch_size, n_id, adjs in subgraph_loader: - # h = features[n_id] - # for l, (edge_index, _, size) in enumerate(adjs): - # # h = layer(self, adj_batch, h_batch, epoch=-1) # GCNConv - # h = self.layers[l](h, edge_index) # SAGEConv - # if l != len(self.layers) - 1: - # h = F.relu(h) - # h = F.log_softmax(h, dim=1) - # batch = n_id[:batch_size] - # print(f"batch: {batch}", flush=True) - # print(f"h[:batch_size]: {h[:batch_size]}", flush=True) - # print(f"h[:batch_size].argmax(dim=-1): {h[:batch_size].argmax(dim=-1)}", flush=True) - # preds = h[:batch_size].argmax(dim=-1) == labels[batch] - # correct_count += preds.sum() - # return correct_count - - # """ SAGEConv inference """ # subgraph_loader = NeighborSampler(graph, node_idx=None, - # sizes=[-1], batch_size=2048, - # shuffle=False) - - # # Compute representations of nodes layer by layer, using *all* - # # available edges. This leads to faster computation in contrast to - # # immediately computing the final representations of each batch. - # total_edges = 0 - # for i in range(self.n_layers): - # xs = [] + # 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, _, size = adj.to(self.device) - # total_edges += edge_index.size(1) - # x = features[n_id].to(self.device) - # # x_target = x[:size[1]] - # # x = self.convs[i]((x, x_target), edge_index) - # x = self.layers[i](x, edge_index) # SAGEConv - # if i != self.n_layers - 1: - # x = F.relu(x) - # x = x[:batch_size] - # xs.append(x) - # features = torch.cat(xs, dim=0) + # 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): @@ -285,7 +260,7 @@ def evaluate(self, graph, features, test_idx, labels): # return features subgraph_loader = NeighborSampler(graph, node_idx=None, - sizes=[-1], batch_size=2048, + sizes=[20], batch_size=1024, # subgraph_loader = NeighborSampler(graph, node_idx=test_idx, # sizes=[-1, -1], batch_size=512, shuffle=False) @@ -668,7 +643,10 @@ def main(args, batches=None): 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) @@ -715,16 +693,30 @@ def main(args, batches=None): if args.dataset == "ogbn-papers100M": print(f"begin copying", flush=True) - adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") - adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") + # adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") + # adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") + 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") print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) print(f"adj_matrix_cols: {adj_matrix_cols} dtype: {adj_matrix_cols.dtype} device: {adj_matrix_cols.device}", flush=True) # adj_matrix_crows = adj_matrix_crows.int().to(device) # adj_matrix_cols = adj_matrix[1,:].int().to(device) + print(f"adj_matrix_crows[104308835]: {adj_matrix_cols[adj_matrix_crows[104308835]:adj_matrix_crows[104308835 + 1]]}", flush=True) adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) - g_loc = torch.sparse_csr_tensor(adj_matrix_crows.int(), adj_matrix_cols.int(), \ + adj_matrix_crows = adj_matrix_crows.int() + adj_matrix_cols = adj_matrix_cols.int() + 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) @@ -734,7 +726,9 @@ def main(args, batches=None): g_loc_crows = g_loc.crow_indices().int() g_loc_cols = g_loc.col_indices().int() g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, adj_matrix_vals, g_loc.size()) - g_loc = g_loc.to(device) + 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 = \ @@ -800,11 +794,11 @@ def main(args, batches=None): 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)) + 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(test_idx_len, dst=dst_gpu) dist.send(train_idx, dst=dst_gpu) - # dist.send(test_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) @@ -901,7 +895,12 @@ def main(args, batches=None): g_loc_values = g_loc_values.cpu() g_loc = torch.sparse_csr_tensor(g_loc_crows, g_loc_cols, g_loc_values, \ torch.Size([g_loc_meta[1], g_loc_meta[2]])) - g_loc = g_loc.to(device) + 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) @@ -952,14 +951,14 @@ def main(args, batches=None): if args.dataset == "ogbn-papers100M": train_idx_len = torch.cuda.LongTensor(1).fill_(0) - # test_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) + 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) + 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) + dist.recv(test_idx, src=src_gpu) torch.cuda.synchronize() features_loc = inputs.to(device) else: @@ -1030,7 +1029,9 @@ def main(args, batches=None): 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 @@ -1047,7 +1048,7 @@ def main(args, batches=None): # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model - # torch.manual_seed(0) + torch.manual_seed(0) if args.sample_method == "sage": model = GCN(num_features, args.n_hidden, @@ -1077,8 +1078,8 @@ def main(args, batches=None): model = model.to(device) # use optimizer - # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) + # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * size) # optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) dur = [] @@ -1283,11 +1284,9 @@ def main(args, batches=None): nnz_row += nnz_row * args.samp_num[j] nnz_col += nnz_col * args.samp_num[j] - print(f"i: {i} nnz_row: {nnz_row} nnz_col: {nnz_col}", flush=True) 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)) - print(f"min: {adj_row_select_min} max: {adj_row_select_max}", flush=True) if epoch >= 1: start_time(start_inner_timer, timing_arg=True) @@ -1352,6 +1351,12 @@ def main(args, batches=None): 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] @@ -1359,10 +1364,15 @@ def main(args, batches=None): 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) @@ -1381,9 +1391,14 @@ def main(args, batches=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) @@ -1573,6 +1588,15 @@ def main(args, batches=None): # adjs = [adj.to(rank) for adj in adjs] # features_batch = features_loc[n_id] optimizer.zero_grad() + print(f"adjs: {adjs}", flush=True) + print(f"features_batch: {features_batch}", flush=True) + print(f"batch_vtxs: {batch_vtxs}", flush=True) + degrees = torch.cuda.LongTensor(batch_vtxs.size(0)) + for vtx_idx, vtx in enumerate(batch_vtxs): + deg = g_loc.crow_indices()[vtx + 1] - g_loc.crow_indices()[vtx] + degrees[vtx_idx] = deg + print(f"degrees: {degrees}", flush=True) + print(f"degrees == 0: {(degrees == 0).sum()}", flush=True) logits = model(adjs, features_batch, epoch) torch.cuda.nvtx.range_pop() # nvtx-fwd if epoch >= 1: @@ -1636,6 +1660,7 @@ def main(args, batches=None): 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]) @@ -1675,20 +1700,48 @@ def main(args, batches=None): # 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): - model = model.cpu() + # 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): train_nid = train_nid.cpu() test_nid = test_nid.cpu() - out = model.evaluate(adj_matrix, inputs, test_nid, data.y) + adj_matrix_csr = g_loc.cpu() + adj_matrix = adj_matrix_csr.clone().detach().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()) - res = out.argmax(dim=-1) == data.y.cpu() - # acc1 = int(res[train_nid].sum()) / train_nid.size(0) - acc3 = int(res[test_nid].sum()) / test_nid.size(0) + 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: + 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) # acc3 = correct_count / test_nid.size(0) - print("Rank: {:05d} | Epoch: {:05d} | Time(s): {:.4f} | Loss: {} | Accuracy: {:.4f}".format(rank, epoch, np.sum(dur), loss.item(), acc3), flush=True) + 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) - model = model.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) diff --git a/examples/pyg_papers.py b/examples/pyg_papers.py new file mode 100644 index 0000000..6daf73b --- /dev/null +++ b/examples/pyg_papers.py @@ -0,0 +1,209 @@ +"""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 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.loader import NeighborLoader +from torch_geometric.nn import GCN, GraphSAGE + + +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, 5], + ) + + print(f"before neighbor loader", flush=True) + train_loader = NeighborLoader( + input_nodes=split_idx['train'], + shuffle=True, + drop_last=True, + **kwargs, + ) + 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) + 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, 7): + model.train() + for i, batch in enumerate(train_loader): + if i == warmup_steps: + torch.cuda.synchronize() + start = time.time() + batch = batch.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] + loss = F.cross_entropy(out, 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): + 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) + 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"] = "nid008460" + 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() + # data.y = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/label/label.pt") + # data.y = data.y.squeeze().to(device) + # 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) + + # 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_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) + # g_loc = torch.sparse_csr_tensor(adj_matrix_crows, adj_matrix_cols, adj_matrix_vals) + # edge_index = g_loc.to_sparse_coo()._indices() + # data.edge_index = edge_index + # print(f"data.edge_index: {data.edge_index}", flush=True) + # print(f"data.edge_index.min: {data.edge_index.min()}", flush=True) + # print(f"data.edge_index.max: {data.edge_index.max()}", flush=True) + # print(f"train_idx: {train_idx}", flush=True) + # print(f"train_idx.min: {train_idx.min()}", flush=True) + # print(f"train_idx.max: {train_idx.max()}", flush=True) + + # inputs = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") + # data.x = inputs + # print(f"inputs: {inputs} inputs.size: {inputs.size()} inputs.dtype {inputs.dtype}", flush=True) + # num_features = inputs.size(1) + # num_classes = 172 + # data.num_features = num_features + + 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, 3, dataset.num_classes) + acc = Accuracy(task="multiclass", num_classes=dataset.num_classes) + # model = GCN(num_features, 256, 2, num_classes) + # acc = Accuracy(task="multiclass", num_classes=num_classes) + data = dataset[0] + data.y = data.y.reshape(-1) + run(nprocs, data, split_idx, model, acc, wall_clock_start) diff --git a/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 93f3207..317ba1b 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2749,20 +2749,25 @@ void rearrangel_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_col } __global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows, int *matb_crows, - int *matb_cols, int *matc_cols, long nnz_count) { + int *matb_cols, int *matc_cols, size_t 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) { + 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]; - long row_len = matb_crows[col + 1] - matb_crows[col]; + int 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] = (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]); + } } } } From 748879b7065b753307a08a8c5a952ddcaec62b47 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Wed, 10 Apr 2024 17:46:57 -0700 Subject: [PATCH 175/178] Miscellaneous optimizations for evaluation --- cagnet/samplers/sage.py | 3 - cagnet/samplers/utils.py | 16 +- examples/gcn_15d.py | 174 +++++++++---- examples/neighbor_sampler.py | 200 +++++++++++++++ examples/pyg_papers.py | 337 +++++++++++++++++++++----- sparse-extension/sparse_coo_tensor.cu | 16 +- 6 files changed, 616 insertions(+), 130 deletions(-) create mode 100644 examples/neighbor_sampler.py diff --git a/cagnet/samplers/sage.py b/cagnet/samplers/sage.py index 8f86ac8..f5adc3e 100644 --- a/cagnet/samplers/sage.py +++ b/cagnet/samplers/sage.py @@ -133,9 +133,6 @@ def sage_sampler(adj_matrix, batches, batch_size, frontier_sizes, mb_count_total timing_dict["sage-startiter"].append(stop_time(start_timer, stop_timer)) # Expand batches matrix - print(f"l: {i} current_frontier: {current_frontier}", flush=True) - print(f"l: {i} current_frontier.indices.min: {current_frontier._indices().min()}", flush=True) - print(f"l: {i} current_frontier.indices.dtype: {current_frontier._indices().dtype}", flush=True) 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", diff --git a/cagnet/samplers/utils.py b/cagnet/samplers/utils.py index d28be2a..dd695bb 100644 --- a/cagnet/samplers/utils.py +++ b/cagnet/samplers/utils.py @@ -1086,28 +1086,18 @@ def gen_prob_dist(numerator, adj_matrix, mb_count, node_count_total, replication # 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] + 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) - print(f"before crow_indices: {adj_matrix.crow_indices()}", flush=True) - print(f"before crow_indices.min: {adj_matrix.crow_indices().min()}", flush=True) - print(f"before crow_indices.max: {adj_matrix.crow_indices().max()}", flush=True) - print(f"before crow_indices.dtype: {adj_matrix.crow_indices().dtype}", flush=True) - print(f"before col_indices: {adj_matrix.col_indices()}", flush=True) - print(f"before col_indices.min: {adj_matrix.col_indices().min()}", flush=True) - print(f"before col_indices.max: {adj_matrix.col_indices().max()}", flush=True) - print(f"before col_indices.dtype: {adj_matrix.col_indices().dtype}", flush=True) - print(f"before matc_chunk_cols.min: {matc_chunk_cols.min()}", flush=True) rearrange_rows_gpu(numerator._indices()[0,:], numerator._indices()[1,:], matc_chunk_crows, - adj_matrix.crow_indices(), + adj_matrix.crow_indices().long(), adj_matrix.col_indices(), matc_chunk_cols) - print(f"after matc_chunk_cols: {matc_chunk_cols}", flush=True) - print(f"after matc_chunk_cols.min: {matc_chunk_cols.min()}", flush=True) # p_num_rows = torch.repeat_interleave( # torch.arange(matc_chunk_crows.size(0) - 1, # dtype=torch.int32, diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 33fa420..916cd7f 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -13,9 +13,10 @@ 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.datasets import Planetoid, Reddit from torch_geometric.loader import NeighborSampler from torch_geometric.utils import add_remaining_self_loops import torch_sparse @@ -78,11 +79,13 @@ def __init__(self, in_feats, n_hidden, n_classes, n_layers, aggr, rank, size, pa # self.layers.append(GCNConv(n_hidden, n_classes, self.partitioning, self.device)) # # # self.layers.append(GCNConv(in_feats, n_classes, self.partitioning, self.device)) - self.layers.append(SAGEConv(in_feats, n_hidden)) - for _ in range(n_layers - 2): - self.layers.append(SAGEConv(n_hidden, n_hidden)) - self.layers.append(SAGEConv(n_hidden, n_classes)) - # self.layers.append(SAGEConv(in_feats, n_classes)) + 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 @@ -108,41 +111,60 @@ def forward(self, graphs, inputs, epoch): def evaluate(self, graph, features, test_idx, labels): # for papers100M - subgraph_loader = NeighborSampler(graph, node_idx=test_idx, + 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] + 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): - edge_index, _, size = adjs[l] + 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) adj_batch = torch.sparse_coo_tensor(edge_index, - torch.FloatTensor(edge_index.size(1)).fill_(1.0), - # torch.cuda.FloatTensor(edge_index.size(1)).fill_(1.0), + # torch.FloatTensor(edge_index.size(1)).fill_(1.0), + values, size) - adj_batch = adj_batch.t().coalesce() + 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]) + 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) + # non_eval_timings = copy.deepcopy(self.timings) # for l, layer in enumerate(self.layers): # hs = [] # print(f"l: {l}", flush=True) @@ -695,18 +717,22 @@ def main(args, batches=None): print(f"begin copying", flush=True) # adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") # adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") - 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_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") - print(f"adj_matrix_crows: {adj_matrix_crows} dtype: {adj_matrix_crows.dtype} device: {adj_matrix_crows.device}", flush=True) + # 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) - print(f"adj_matrix_crows[104308835]: {adj_matrix_cols[adj_matrix_crows[104308835]:adj_matrix_crows[104308835 + 1]]}", flush=True) adj_matrix_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) - adj_matrix_crows = adj_matrix_crows.int() + # 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)])) @@ -849,7 +875,8 @@ def main(args, batches=None): # 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.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() print(f"recv crows", flush=True) @@ -1004,6 +1031,7 @@ def main(args, batches=None): 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) dist.barrier() @@ -1048,7 +1076,7 @@ def main(args, batches=None): # return frontiers, adj_matrices, adj_matrix, col_groups # create GCN model - torch.manual_seed(0) + # torch.manual_seed(0) if args.sample_method == "sage": model = GCN(num_features, args.n_hidden, @@ -1078,7 +1106,13 @@ def main(args, batches=None): model = model.to(device) # use optimizer - optimizer = torch.optim.Adam(model.parameters(), lr=args.lr * math.sqrt(size)) + 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) @@ -1217,6 +1251,55 @@ def main(args, batches=None): 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)) @@ -1364,15 +1447,15 @@ def main(args, batches=None): 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) + # 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) + # 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) @@ -1391,14 +1474,14 @@ def main(args, batches=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"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) + # 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) @@ -1588,15 +1671,6 @@ def main(args, batches=None): # adjs = [adj.to(rank) for adj in adjs] # features_batch = features_loc[n_id] optimizer.zero_grad() - print(f"adjs: {adjs}", flush=True) - print(f"features_batch: {features_batch}", flush=True) - print(f"batch_vtxs: {batch_vtxs}", flush=True) - degrees = torch.cuda.LongTensor(batch_vtxs.size(0)) - for vtx_idx, vtx in enumerate(batch_vtxs): - deg = g_loc.crow_indices()[vtx + 1] - g_loc.crow_indices()[vtx] - degrees[vtx_idx] = deg - print(f"degrees: {degrees}", flush=True) - print(f"degrees == 0: {(degrees == 0).sum()}", flush=True) logits = model(adjs, features_batch, epoch) torch.cuda.nvtx.range_pop() # nvtx-fwd if epoch >= 1: @@ -1607,7 +1681,8 @@ def main(args, batches=None): 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.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 @@ -1702,18 +1777,20 @@ def main(args, batches=None): 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): - train_nid = train_nid.cpu() + # 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.clone().detach().to_sparse_coo()._indices() + 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() + if args.dataset == "ogbn-papers100M" or args.dataset == "cora": + # 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) @@ -1727,8 +1804,10 @@ def main(args, batches=None): acc3 = int(res.sum()) / test_nid.size(0) del out del res - model = model.to(device) + # 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) @@ -1737,6 +1816,7 @@ def main(args, batches=None): 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) 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 index 6daf73b..97d7902 100644 --- a/examples/pyg_papers.py +++ b/examples/pyg_papers.py @@ -6,6 +6,7 @@ --container-mounts=/ogb-papers100m/:/workspace/dataset python3 path_to_script.py """ +import gc import os import time from typing import Optional @@ -17,10 +18,185 @@ 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 -from torch_geometric.nn import GCN, GraphSAGE +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 @@ -57,23 +233,35 @@ def run(world_size, data, split_idx, model, acc, wall_clock_start): 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, 5], - ) + # 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, - ) + # 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 = 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 @@ -88,17 +276,24 @@ def run(world_size, data, split_idx, model, acc, wall_clock_start): print("Beginning training...") print(f"len(train_loader): {len(train_loader)}", flush=True) - for epoch in range(1, 7): + for epoch in range(1, 6): model.train() - for i, batch in enumerate(train_loader): + # 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) + # 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] - loss = F.cross_entropy(out, y) + # 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() @@ -114,12 +309,19 @@ def run(world_size, data, split_idx, model, acc, wall_clock_start): @torch.no_grad() def test(loader: NeighborLoader, num_steps: Optional[int] = None): model.eval() - for j, batch in enumerate(loader): + # 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) + # 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 @@ -155,55 +357,70 @@ def test(loader: NeighborLoader, num_steps: Optional[int] = None): if "SLURM_NTASKS" in os.environ.keys(): os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] - os.environ["MASTER_ADDR"] = "nid008460" + 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") + # dataset = PygNodePropPredDataset(name='ogbn-papers100M', root="/global/u1/a/alokt/data") device = torch.device(f'cuda:{rank % 1}') - # data = Data() - # data.y = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/label/label.pt") - # data.y = data.y.squeeze().to(device) - # 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) - - # 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_vals = torch.FloatTensor(adj_matrix_cols.size(0)).fill_(1) - # g_loc = torch.sparse_csr_tensor(adj_matrix_crows, adj_matrix_cols, adj_matrix_vals) - # edge_index = g_loc.to_sparse_coo()._indices() - # data.edge_index = edge_index - # print(f"data.edge_index: {data.edge_index}", flush=True) - # print(f"data.edge_index.min: {data.edge_index.min()}", flush=True) - # print(f"data.edge_index.max: {data.edge_index.max()}", flush=True) - # print(f"train_idx: {train_idx}", flush=True) - # print(f"train_idx.min: {train_idx.min()}", flush=True) - # print(f"train_idx.max: {train_idx.max()}", flush=True) - - # inputs = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/feat/feature.pt") - # data.x = inputs - # print(f"inputs: {inputs} inputs.size: {inputs.size()} inputs.dtype {inputs.dtype}", flush=True) - # num_features = inputs.size(1) - # num_classes = 172 - # data.num_features = num_features - - split_idx = dataset.get_idx_split() + 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, 3, dataset.num_classes) - acc = Accuracy(task="multiclass", num_classes=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] + # 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/sparse-extension/sparse_coo_tensor.cu b/sparse-extension/sparse_coo_tensor.cu index 317ba1b..fe9b852 100644 --- a/sparse-extension/sparse_coo_tensor.cu +++ b/sparse-extension/sparse_coo_tensor.cu @@ -2748,7 +2748,8 @@ void rearrangel_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_col 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, 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; @@ -2758,16 +2759,16 @@ __global__ void RearrangeRows(long *mata_rows, long *mata_cols, long *matc_crows long row = mata_rows[i]; long col = mata_cols[i]; long dst_idx = matc_crows[row]; - int src_idx = matb_crows[col]; + 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]); - } + // 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]); + // } } } } @@ -2788,7 +2789,8 @@ void rearrange_rows_gpu(const at::Tensor& mata_rows, const at::Tensor& mata_cols RearrangeRows<<>>(mata_rows.data(), mata_cols.data(), matc_crows.data(), - matb_crows.data(), + // matb_crows.data(), + matb_crows.data(), matb_cols.data(), matc_cols.data(), mata_rows.sizes()[0]); From 5d537ab037a1cb7b54eb4e1db88718e409ff10cc Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Fri, 24 May 2024 12:25:24 -0700 Subject: [PATCH 176/178] cpu2gpu microbenchmark --- .gitignore | 1 + tests/test_nccl_p2p.py | 211 ++++++++++++++++++++++------------------- 2 files changed, 116 insertions(+), 96 deletions(-) diff --git a/.gitignore b/.gitignore index 666d777..0d499ca 100644 --- a/.gitignore +++ b/.gitignore @@ -39,3 +39,4 @@ examples/*.batch *.nsys-rep *.ncu-rep examples/wandb +*.png diff --git a/tests/test_nccl_p2p.py b/tests/test_nccl_p2p.py index 1abe611..e9546a1 100644 --- a/tests/test_nccl_p2p.py +++ b/tests/test_nccl_p2p.py @@ -26,6 +26,7 @@ def test_nccl(args): rank = dist.get_rank() size = dist.get_world_size() 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) @@ -355,102 +356,120 @@ def test_nccl(args): # torch.cuda.nvtx.range_pop() # torch.cuda.profiler.cudart().cudaProfilerStop() - print(f"alltoall", flush=True) - buffer_size = 1 * 2**30 // size - run_count = 2 - for r in range(run_count): - tensors = [] - for i in range(size): - tensors.append(torch.cuda.FloatTensor(buffer_size).fill_(i)) - tensors[-1] = torch.cuda.FloatTensor(0).fill_(i) - - recv_tensors = [] - for i in range(size): - if rank == size - 1: - recv_tensors.append(torch.cuda.FloatTensor(0)) - else: - recv_tensors.append(torch.cuda.FloatTensor(buffer_size)) - - if r == run_count - 1 and rank == 0: - start_time(start_timer) - - # dist.scatter(recv_tensor, tensors, src=0) - dist.all_to_all(recv_tensors, tensors) - torch.cuda.synchronize() - - if r == run_count - 1 and rank == 0: - seconds = stop_time(start_timer, stop_timer) / 1000 - gb_count = (buffer_size * 4 * size) / 2**30 - bw = gb_count / seconds - print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") - print(f"time(ms): {seconds * 1000}") - dist.barrier() - - dist.barrier() - print(f"batched isend/irecv alltoall", flush=True) - buffer_size = 1 * 2**30 // size - # buffer_size = ((66 // 4) * 2**20 // 2) // (size - 1) - run_count = 3 - for r in range(run_count): - send_tensors = [None] * size - recv_tensors = [None] * size - elem_count = 0 - for i in range(size): - if i != rank: - elem_count += 2 * buffer_size - send_tensors[i] = torch.cuda.FloatTensor(buffer_size).fill_(i + 1) - recv_tensors[i] = torch.cuda.FloatTensor(buffer_size).fill_(0) - - ops = [] - for i in range(size): - if i != rank: - ops.append(dist.P2POp(dist.isend, send_tensors[i], i)) - ops.append(dist.P2POp(dist.irecv, recv_tensors[i], i)) - - if r == run_count - 1: - start_time(start_timer) - reqs = dist.batch_isend_irecv(ops) - for req in reqs: - req.wait() - torch.cuda.synchronize() - dist.barrier() - - if r == run_count - 1 and rank == 0: - seconds = stop_time(start_timer, stop_timer) / 1000 - gb_count = (buffer_size * size * 4) / 2**30 # assumes buffer_size is per process - bw = gb_count / seconds - print(f"gb: {gb_count} GB time: {seconds}s bw: {bw}GB/s") - print(f"time(ms): {seconds * 1000}") - print(f"elem_count: {elem_count}") - - print(f"p2p cross-node send", flush=True) - buffer_size = 500 * 2**20 - # buffer_size = ((66 // 4) * 2**20 // 2) // (size - 1) - 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"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"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') From bd853d1ac52b825e38f19d51b9cddfb427e1bf98 Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Tue, 4 Jun 2024 00:59:31 -0700 Subject: [PATCH 177/178] Tweaks to get papers100M to run again --- examples/gcn_15d.py | 45 +++++++++++++++++++++++---------------------- 1 file changed, 23 insertions(+), 22 deletions(-) diff --git a/examples/gcn_15d.py b/examples/gcn_15d.py index 916cd7f..e4f6ebf 100644 --- a/examples/gcn_15d.py +++ b/examples/gcn_15d.py @@ -115,12 +115,12 @@ def evaluate(self, graph, features, test_idx, labels): 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, + # 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) @@ -138,8 +138,9 @@ def evaluate(self, graph, features, test_idx, labels): # 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())) + # 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, @@ -715,22 +716,22 @@ def main(args, batches=None): if args.dataset == "ogbn-papers100M": print(f"begin copying", flush=True) - # adj_matrix_crows = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indptr.pt") - # adj_matrix_cols = torch.load("/global/u1/a/alokt/data/ogbn_papers100M/csr/indices.pt") + 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") + # 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_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, \ @@ -875,16 +876,14 @@ def main(args, batches=None): # 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) + 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() - print(f"recv crows", flush=True) g_loc_cols = torch.cuda.IntTensor(g_loc_meta[0].item()) dist.recv(g_loc_cols, src=src_gpu) torch.cuda.synchronize() - print(f"recv cols", flush=True) # g_loc_values = torch.cuda.DoubleTensor(g_loc_meta[0].item()) # dist.recv(g_loc_values, src=src_gpu) @@ -899,16 +898,13 @@ def main(args, batches=None): inputs = torch.cuda.FloatTensor(rows, num_features) dist.recv(inputs, src=src_gpu) torch.cuda.synchronize() - print(f"recv inputs", flush=True) features_zero_row = torch.cuda.FloatTensor(num_features) dist.recv(features_zero_row, src=src_gpu) torch.cuda.synchronize() - print(f"recv features_zero_row", flush=True) # data = Data() data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) - print(f"data.y.size: {data.y.size()}", flush=True) dist.recv(data.y, src=src_gpu) torch.cuda.synchronize() data.y = data.y.long() @@ -952,7 +948,6 @@ def main(args, batches=None): features_zero_row = torch.cuda.FloatTensor(num_features) dist.recv(features_zero_row, src=src_gpu) torch.cuda.synchronize() - print(f"recv features_zero_row", flush=True) # data = Data() data.y = torch.cuda.FloatTensor(g_loc_meta[2].item()) @@ -1033,7 +1028,9 @@ def main(args, batches=None): 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() @@ -1186,7 +1183,11 @@ def main(args, batches=None): args.n_bulkmb = 1 args.batch_size = batch_count - b torch.cuda.nvtx.range_push("nvtx-partition-batches") - batches = batches_all[b:(b + args.n_bulkmb * args.batch_size)].view(args.n_bulkmb, args.batch_size) + + 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": @@ -1789,7 +1790,7 @@ def main(args, batches=None): 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" or args.dataset == "cora": + 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) From f820a567a8b367fe0c9e01fe40cb2e71a68a3a6b Mon Sep 17 00:00:00 2001 From: Alok Tripathy Date: Thu, 8 Aug 2024 12:10:57 -0700 Subject: [PATCH 178/178] Updated README --- README.md | 87 +++++++++++++++++++++++-------------------------------- 1 file changed, 37 insertions(+), 50 deletions(-) diff --git a/README.md b/README.md index 33f8685..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.8.11 -- PyTorch 1.9.0 -- PyTorch Geometric (PyG) 1.7.2 -- CUDA 10.2 -- GCC 9.3.0 +- Python 3.9.19 +- PyTorch 1.13 +- PyTorch Geometric (PyG) 2.5.3 +- CUDA 11.7 +- GCC 11.2 -On NERSC Perlmutter, 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 tudatoolkit/21.3_10.2 # CUDA 10.2 -module load pytorch/1.9.0 # PyTorch and PyG -``` +module load pytorch/1.13 -On OLCF Summit, older versions of these dependencies can be accessed with the following -```bash -module load cuda # CUDA 10.1 -module load gcc # GCC 6.4.0 -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 ``` - ## 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. -- `examples/gcn_1d.py` : 1D algorithm +Our minibatch training pipeline, including both our SAGE and LADIES implementations, can be accessed by running - `examples/gcn_15d.py` : 1.5D algorithm -- `examples/gcn_2d.py` : 2D algorithm -- `examples/gcn_3d.py` : 3D algorithm -Each file also as the following flags: +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 +- `--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 -- `--normalize` : Normalize adjacency matrix in preprocessing -- `--accuracy ` : Compute and print accuracy metrics (Reddit only) -- `--replication ` : Replication factor (1.5D algorithm only) - -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 Cora and Reddit, PyG handles downloading and accessing the dataset. +- `--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 Reddit with -- 4 processes -- 100 epochs -- 16 hidden layer activations -- 2-factor replication -- normalization -- no weight decay +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 -n 4 --ntasks-per-node=4 --gpus-per-task=1 --gpu-bind=map_gpu:0,1,2,3 python gcn_15d.py --dataset=Reddit --n-epochs=100 --n-hidden 16 --weight-decay 0 --replication 2 --normalize` +`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.