This repository has been archived by the owner on Oct 21, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 3
/
util.py
204 lines (172 loc) · 6.79 KB
/
util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
import dgl
import numpy as np
import random
import networkx as nx
import torch
import json
import matplotlib.pyplot as plt
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)
ROW = 2
COL = 2
PIPE_CYCLE = 4
def create_graph(graphdef, numnodes = 10):
# random generate a directed acyclic graph
if graphdef is None:
graphdef = {}
a = nx.generators.directed.gn_graph(numnodes)
# a = nx.generators.directed.gnc_graph(numnodes)
graph = dgl.from_networkx(a)
# make one input node
asc = dgl.topological_nodes_generator(graph)
for i in range(1, len(asc[0])):
graph.add_edge(asc[0][0], asc[0][i])
tm_idx_total = random.randint(1, numnodes//2) #number of tm variables
tile_memory_req = {}
for i in range(0, tm_idx_total):
nodeid = random.randint(0, numnodes-1)
if nodeid in tile_memory_req:
tile_memory_req[nodeid].append(i)
else:
tile_memory_req[nodeid] = [i]
graphdef['nodes_to_tm'] = tile_memory_req
else:
tile_memory_req = graphdef['nodes_to_tm']
edges = graphdef['graphdef']
graph = dgl.graph((torch.Tensor(edges[0]).int(), torch.Tensor(edges[1]).int()))
if len(edges) == 3 and edges[2] > 0:
graph.add_nodes(edges[2])
tm_idx_total = len(graphdef['tile_memory_map'].keys())
# Create TM to node mappings
tm_to_nodes = {tm_idx: [] for tm_idx in range(0, tm_idx_total)}
for instr_idx, tm_idxs in tile_memory_req.items():
for tm_idx in tm_idxs:
tm_to_nodes[tm_idx].append(instr_idx)
graphdef['tm_to_nodes'] = tm_to_nodes
# Add tile memory constraints as features to graph
tm_req_feat = torch.zeros(graph.num_nodes(), tm_idx_total) # node, tile mem var index binary vector
for instr_idx, tm_idxs in tile_memory_req.items():
for tm_idx in tm_idxs:
tm_req_feat[instr_idx][tm_idx] = 1
graph.ndata['tm_req'] = tm_req_feat
# create list of sync start flow nodes
sf_nodes = []
for node in graph.nodes():
if len(graph.predecessors(node).numpy()) == 0:
sf_nodes.append(node.item())
graphdef['graph'] = graph
graphdef['sf_nodes'] = sf_nodes
return graphdef
def positional_encoding(pos, feat_size=16, timescale=10000):
'''
pos : [N X D] matrix of positions. N is the number of slices.
returns a positional encoding of [N x (D * feat_size)]
'''
N, D = pos.shape
sin_freq = torch.arange(0, feat_size, 2.0) / feat_size
cos_freq = torch.arange(1, feat_size, 2.0) / feat_size
sin_freq = 1 / (timescale ** sin_freq)
cos_freq = 1 / (timescale ** cos_freq)
sin_emb = torch.sin(torch.einsum('ni,d->ndi', pos, sin_freq))
cos_emb = torch.cos(torch.einsum('ni,d->ndi', pos, cos_freq))
encoding = torch.zeros(N, D * feat_size)
for i in range(D):
start_idx = i * feat_size
end_idx = (i + 1) * feat_size
encoding[:, start_idx:end_idx:2] = sin_emb[:, :, i]
encoding[:, start_idx+1:end_idx:2] = cos_emb[:, :, i]
return encoding
def get_graph_json(path):
with open(path) as file: # Use file to refer to the file object
data = json.load(file)
edge_src = []
edge_dst = []
tmem_map = {} # tile_mem variable str : int index
nidx = 0
for mem in data['TileMemories'].keys(): # give index to each tile mem variable
tmem_map[mem] = nidx
nidx += 1
nidx = 0
for graph in data['Program']: # graphs
offset = nidx
for node in graph['SyncFlow']: # nodes
for edges in node['SEInst']['Successors']:
edge_src.append(nidx)
edge_dst.append(edges + offset)
nidx += 1
extra_node = (nidx-1) - max(max(edge_src), max(edge_dst))
nidx = 0
tmem_req = {} # node id : list of tile mem var indexes
for graph in data['Program']: # graphs
for node in graph['SyncFlow']: # nodes
l = [] #tile mem var indexes
for var in node['SEInst']['SEInstUse']:
if var in tmem_map:
l.append(tmem_map[var])
tmem_req[nidx] = l
nidx += 1
return {'graphdef': (edge_src, edge_dst, extra_node),
'nodes_to_tm': tmem_req,
'tile_memory_map': tmem_map}
def print_graph(graphdef):
graph_in = graphdef['graph'].adjacency_matrix_scipy().toarray()
print('graph adjacency matrix: ', graph_in)
nx_g = graphdef['graph'].to_networkx()
nx.draw(nx_g, nx.nx_agraph.graphviz_layout(nx_g, prog='dot'), with_labels=True)
plt.show()
def output_json(placed_nodes, no_of_tiles=16, spoke_count=3 ,out_file_name='mapping.json'):
"""[summary]
Args:
instr_coords (np.array): Array w/ shape [Number of slices, 3]
out_file (str, optional): Output json name. Defaults to 'mapping.json'.
"""
data = {}
#TODO: Change when using variable spoke count for each tile
num_spokes = [spoke_count for _ in range(no_of_tiles)]
data['num_tiles'] = no_of_tiles
data['num_spokes'] = num_spokes
mappings = [{'tile_id': tile_idx, 'spoke_map': ['' for _ in range(spoke_count)]} for \
tile_idx in range (no_of_tiles)]
# Iterate over assignment
"""
for instr_idx, tile_coord in enumerate(instr_coords):
tile_idx = int(tile_coord[0])
spoke_no = int(tile_coord[2])
mappings[tile_idx]['spoke_map'][spoke_no] = f'instruction ID#{instr_idx}'
"""
for node_idx, info in placed_nodes.items():
tile_idx = info['tile_slice'][0]
spoke_no = info['tile_slice'][1]
mappings[tile_idx]['spoke_map'][spoke_no] = f'instruction ID#{node_idx}'
data['mappings'] = mappings
with open(out_file_name, 'w') as outfile:
json.dump(data, outfile, indent=4)
def ravel_index(pos, shape):
res = 0
acc = 1
for pi, si in zip(reversed(pos), reversed(shape)):
res += pi * acc
acc *= si
return res
def initial_fill(num_nodes, grid_shape, manual = None):
'''
fill array random:
return
grid: grid with node number in the array
grid_in: list of node and its placed coord
manual: list of index coord for each node
'''
grid = -np.ones(grid_shape)# -1 is unassigned
grid_in = []
gg = np.prod(grid_shape)
if isinstance(manual, list):
place = manual
else:
place = random.sample(range(gg), num_nodes) # list of unique elements chosen from the population sequence
for i, idx in enumerate(place):
x, y, c = np.unravel_index(idx, grid_shape) # index to [coord c, y, x]
grid[x, y, c] = i
grid_in.append([x, y, c])
grid_in = np.array(grid_in)
return grid, grid_in, place