1+ "New functions created by us. Only for illustration, cannot be executed like this in this script. They are copied out of the partitioner.py line 624 to 779"
2+
3+ def lr_boundary_seq (self , rank : int ):
4+ """Calculates the sequence of boundaries on an edge of a tile + the amount of to_ranks shared with this rank"""
5+ rank_position_y = (rank % self .tile .total_ranks ) // self .tile .layout [0 ] # position of rank on left/right boundary in y-axis
6+ boundary_seq = []
7+ if self .tile .layout [0 ] % self .tile .layout [1 ] == 0 or self .tile .layout [1 ] % self .tile .layout [0 ] == 0 : # layout is dividable
8+ if self .tile .layout [0 ] >= self .tile .layout [1 ]:
9+ div = self .tile .layout [0 ] // self .tile .layout [1 ]
10+ else :
11+ div = 1
12+ boundary_seq = [div ]* self .tile .layout [1 ] # adds sequence of repeated multiples (div) in a list.
13+ Nboundary = boundary_seq [rank_position_y ]
14+ else :
15+ nfirst = self .tile .layout [0 ] // self .tile .layout [1 ] + 1
16+ boundary_seq .append (nfirst )
17+ for y in range (1 ,self .tile .layout [1 ]): # calculates the rest of the sequence of the boundary by dividing the layout
18+ num1 = (self .tile .layout [0 ]* self .tile .layout [1 ]) - y * self .tile .layout [0 ]
19+ num2before = num1 // self .tile .layout [1 ] + 1
20+ num2after = (num1 - self .tile .layout [0 ])// self .tile .layout [1 ] + 1
21+ numadd = num2before - num2after + 1
22+ boundary_seq .append (numadd )
23+ Nboundary = boundary_seq [rank_position_y ] # Amount of boundaries created on the edge of a rank
24+ return Nboundary , boundary_seq
25+
26+ def ul_boundary_seq (self , rank : int ):
27+ """Calculates the sequence of boundaries on an edge of a tile + the amount of to_ranks shared with this rank"""
28+ rank_position_x = (rank % self .tile .total_ranks )// self .tile .layout [0 ] # position of rank on upper/ower boundary in x-axis
29+ boundary_seq = []
30+ if self .tile .layout [0 ] % self .tile .layout [1 ] == 0 or self .tile .layout [1 ] % self .tile .layout [0 ] == 0 : # layout is dividable
31+ if self .tile .layout [1 ] >= self .tile .layout [0 ]:
32+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
33+ else :
34+ div = 1
35+ boundary_seq = [div ]* self .tile .layout [0 ]
36+ Nboundary = boundary_seq [rank_position_x ]
37+ else :
38+ nfirst = self .tile .layout [1 ] // self .tile .layout [0 ] + 1
39+ boundary_seq .append (nfirst )
40+ for x in range (1 ,self .tile .layout [0 ]): # calculates the rest of the sequence of the boundary by dividing the layout
41+ num1 = (self .tile .layout [0 ]* self .tile .layout [1 ]) - x * self .tile .layout [1 ]
42+ num2before = num1 // self .tile .layout [0 ] + 1
43+ num2after = (num1 - self .tile .layout [1 ])// self .tile .layout [0 ] + 1
44+ numadd = num2before - num2after + 1
45+ boundary_seq .append (numadd )
46+ Nboundary = boundary_seq [rank_position_x ] # Amount of boundaries created on the edge of a rank
47+ return Nboundary , boundary_seq
48+
49+ def lr_to_ranks (self , rank : int , to_root_rank : int ):
50+ """Returns the to_ranks of the boundaries for a specific rank and its edge"""
51+ Nboundary , boundary_seq = self .lr_boundary_seq (rank )
52+ if self .tile .layout [0 ] % self .tile .layout [1 ] == 0 or self .tile .layout [1 ] % self .tile .layout [0 ] == 0 :
53+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
54+ boundary_seq = boundary_seq
55+ if (div >= 1 ):
56+ cumsum_boundary_seq = np .cumsum (boundary_seq )// div # cumulated sum of boundary sequence to calculate shared ranks later on
57+ else :
58+ cumsum_boundary_seq = np .cumsum (boundary_seq )
59+ else :
60+ boundary_seq = [x - 1 for x in boundary_seq ]
61+ cumsum_boundary_seq = np .cumsum (boundary_seq ) # a cumulative sum of boundary_seq to know the start and end of the shared ranks
62+ rank_position_y = (rank % self .tile .total_ranks ) // self .tile .layout [0 ]
63+ to_root_rank = to_root_rank
64+ to_ranks_pot = []
65+ for x in range (self .tile .layout [0 ]):
66+ if self .tile .on_tile_left (rank ):
67+ to_ranks_pot .append (to_root_rank + (self .tile .layout [0 ]* (self .tile .layout [1 ]- 1 )) + x )
68+ else :
69+ to_ranks_pot .append (to_root_rank + x ) # creates list of all sharable ranks (potential to_ranks) on the respective edge
70+ to_ranks_pot = np .fliplr ([to_ranks_pot , to_ranks_pot , to_ranks_pot ])[1 ] # flips list for orientation purposes
71+ if (rank_position_y == 0 ):
72+ if self .tile .layout [0 ] % self .tile .layout [1 ] == 0 or self .tile .layout [1 ] % self .tile .layout [0 ] == 0 :
73+ if self .tile .layout [0 ]// self .tile .layout [1 ] >= 1 :
74+ start = 0
75+ end = cumsum_boundary_seq [0 ]
76+ to_ranks_seq = to_ranks_pot [start :end ]
77+ else :
78+ start = 0
79+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
80+ end = cumsum_boundary_seq [0 ]// div
81+ to_ranks_seq = to_ranks_pot [start :end ]
82+ else :
83+ start = 0
84+ end = cumsum_boundary_seq [0 ]+ 1
85+ to_ranks_seq = to_ranks_pot [start :end ]
86+ else :
87+ if self .tile .layout [0 ]% self .tile .layout [1 ] == 0 or self .tile .layout [1 ]% self .tile .layout [0 ] == 0 :
88+ if self .tile .layout [0 ]/ self .tile .layout [1 ] >= 1 :
89+ start = cumsum_boundary_seq [rank_position_y - 1 ]
90+ end = cumsum_boundary_seq [rank_position_y ]
91+ to_ranks_seq = to_ranks_pot [start :end ]
92+ else :
93+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
94+ start = rank_position_y // div
95+ end = start + 1
96+ to_ranks_seq = to_ranks_pot [start :end ]
97+ else :
98+ start = cumsum_boundary_seq [rank_position_y - 1 ]
99+ end = cumsum_boundary_seq [rank_position_y ] + 1
100+ to_ranks_seq = to_ranks_pot [start :end ]
101+ return to_ranks_seq
102+
103+ def ul_to_ranks (self , rank : int , to_root_rank : int ):
104+ """Returns the to_ranks of the boundaries for a specific rank and its edge"""
105+ Nboundary , boundary_seq = self .ul_boundary_seq (rank )
106+ if self .tile .layout [0 ]% self .tile .layout [1 ] == 0 or self .tile .layout [1 ]% self .tile .layout [0 ] == 0 :
107+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
108+ boundary_seq = boundary_seq
109+ if (div >= 1 ):
110+ cumsum_boundary_seq = np .cumsum (boundary_seq )// div # cumulated sum of boundary sequence to calculate shared ranks later on
111+ else :
112+ cumsum_boundary_seq = np .cumsum (boundary_seq )
113+ else :
114+ boundary_seq = [x - 1 for x in boundary_seq ]
115+ cumsum_boundary_seq = np .cumsum (boundary_seq ) # a cumulative sum of boundary_seq to know the start and end of the shared ranks
116+ rank_position_x = (rank % self .tile .total_ranks )% self .tile .layout [0 ]
117+ to_root_rank = to_root_rank
118+ to_ranks_pot = []
119+ for y in range (self .tile .layout [1 ]):
120+ if self .tile .on_tile_top (rank ):
121+ to_ranks_pot .append (to_root_rank + self .tile .layout [0 ]* y ) # creates list of all sharable ranks (potential to_ranks) on the respective edge
122+ else :
123+ to_ranks_pot .append (to_root_rank + self .tile .layout [0 ]* y + self .tile .layout [0 ]- 1 )
124+
125+ to_ranks_pot = np .fliplr ([to_ranks_pot ,to_ranks_pot ,to_ranks_pot ])[1 ] # flips list for orientation
126+ if (rank_position_x == 0 ):
127+ if self .tile .layout [0 ]% self .tile .layout [1 ] == 0 or self .tile .layout [1 ]% self .tile .layout [0 ] == 0 :
128+ if self .tile .layout [1 ]// self .tile .layout [0 ] > 1 :
129+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
130+ start = 0
131+ end = start + div
132+ to_ranks_seq = to_ranks_pot [start :end ]
133+ else :
134+ div = self .tile .layout [0 ]// self .tile .layout [1 ]
135+ start = 0
136+ end = start + 1
137+ to_ranks_seq = to_ranks_pot [start :end ]
138+ else :
139+ start = 0
140+ end = cumsum_boundary_seq [0 ]+ 1
141+ to_ranks_seq = to_ranks_pot [start :end ]
142+ else :
143+ if self .tile .layout [0 ]% self .tile .layout [1 ] == 0 or self .tile .layout [1 ]% self .tile .layout [0 ] == 0 :
144+ if self .tile .layout [1 ]// self .tile .layout [0 ] > 1 :
145+ div = self .tile .layout [1 ]// self .tile .layout [0 ]
146+ start = rank_position_x * div
147+ end = start + div
148+ to_ranks_seq = to_ranks_pot [start :end ]
149+ else :
150+ div = self .tile .layout [0 ]// self .tile .layout [1 ]
151+ start = rank_position_x // div
152+ end = start + 1
153+ to_ranks_seq = to_ranks_pot [start :end ]
154+ else :
155+ start = cumsum_boundary_seq [rank_position_x - 1 ]
156+ end = cumsum_boundary_seq [rank_position_x ] + 1
157+ to_ranks_seq = to_ranks_pot [start :end ]
158+ return to_ranks_seq
0 commit comments