torch.cuda.synchronize(device=device)

print(f"Sleeping rank {rank}", flush=True)

time.sleep(20)

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)

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()

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()

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)

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)

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)

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)

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))))

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

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:]

@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)

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()