1+
2+ import os
3+ import sys
4+ import sst
5+ from sst .merlin .base import *
6+ from sst .merlin .endpoint import *
7+ from sst .merlin .interface import *
8+ from sst .merlin .topology import *
9+ from sst .ember import *
10+
11+ try :
12+ import networkx as nx
13+ except ImportError :
14+ print ("NetworkX is required. Please install NetworkX and try again." )
15+ sys .exit (1 )
16+
17+
18+ def prepare_complete_4 (write_to_file : bool = False ):
19+ '''
20+ Prepare a complete graph with 4 vertices.\n
21+ If write_to_file is True, write the graph to a GraphML file.\n
22+ Returns the graph (or file path) and topology name.
23+ '''
24+ # Create a complete graph with 4 vertices
25+ G = nx .complete_graph (4 )
26+ # Add endpoint attributes to each vertex
27+ # For this test, we'll assign 2 endpoints per vertex
28+ for node in G .nodes ():
29+ G .nodes [node ]['endpoints' ] = 2
30+
31+ if write_to_file :
32+ # Write the graph to a GraphML file
33+ graph_file = "test_complete_4.graphml"
34+ nx .write_graphml (G , graph_file )
35+ print (f"Generated complete graph with 4 vertices and saved to { graph_file } )
36+ return graph_file , "complete_4"
37+ else :
38+ return G , "complete_4"
39+
40+ def prepare_cubical_graph (write_to_file : bool = False ):
41+ '''
42+ Prepare a cubical graph with 8 vertices and 12 edges.\n
43+ If write_to_file is True, write the graph to a GraphML file.\n
44+ Returns the graph (or file path) and topology name.
45+ '''
46+ # Create a cubical graph with 8 vertices
47+ G = nx .cubical_graph ()
48+ # Add endpoint attributes to each vertex
49+ # For this test, we'll assign 2 endpoints per vertex
50+ for node in G .nodes ():
51+ G .nodes [node ]['endpoints' ] = 2
52+
53+ if write_to_file :
54+ # Write the graph to a GraphML file
55+ graph_file = "test_cubical.graphml"
56+ nx .write_graphml (G , graph_file )
57+ print (f"Generated cubical graph with 8 vertices and saved to { graph_file } )
58+ return graph_file , "cubical"
59+ else :
60+ return G , "cubical"
61+
62+ def run_sst (input_graph : str | nx .Graph , topo_name : str , SR_routing_table : dict = None ):
63+
64+ UNIFIED_ROUTER_LINK_BW = UNIFIED_ROUTER_LINK_BW = 16
65+
66+ BENCH = "FFT3D"
67+ BENCH_PARAMS = " nx=100 ny=100 nz=100 npRow=2"
68+ CPE = 1
69+ EPR = input_graph .nodes [0 ]['endpoints' ] # Endpoints per router
70+
71+ ### Setup the topology, read graph from file
72+ topo = topoAny ()
73+ topo .routing_mode = "source_routing"
74+ topo .topo_name = topo_name
75+ topo .import_graph (input_graph )
76+ # topo.set_verbose_level(16) # Set verbose level for debugging
77+
78+ defaults_z = PlatformDefinition .compose ("firefly-defaults-Z" ,[("firefly-defaults" ,"ALL" )])
79+ defaults_z .addParamSet ("nic" ,{
80+ "num_vNics" : CPE ,
81+ })
82+ PlatformDefinition .setCurrentPlatform ("firefly-defaults-Z" )
83+ # PlatformDefinition.setCurrentPlatform("firefly-defaults")
84+
85+
86+ # Set up the routers
87+ router = hr_router ()
88+ router .link_bw = f"{ UNIFIED_ROUTER_LINK_BW }
89+ router .flit_size = "32B"
90+ router .xbar_bw = f"{ UNIFIED_ROUTER_LINK_BW * 2 } # 2x crossbar speedup
91+ router .input_latency = "20ns"
92+ router .output_latency = "20ns"
93+ router .input_buf_size = "32kB"
94+ router .output_buf_size = "32kB"
95+ router .num_vns = 2
96+ router .xbar_arb = "merlin.xbar_arb_rr"
97+
98+ topo .router = router
99+ topo .link_latency = "20ns"
100+ topo .host_link_latency = "10ns"
101+
102+ _SR_routing_table : dict
103+
104+ if SR_routing_table is not None :
105+ # The source routing table can be passed from input arg
106+ _SR_routing_table = SR_routing_table
107+ else :
108+ # Otherwise, the source routing tables can be calculated here from the topology class method
109+ # by default, use All Shortest Paths (ASP) will be calculated
110+ _SR_routing_table = topo .calculate_routing_table ()
111+
112+ ### use endpointNIC
113+ endpointNIC = EndpointNIC (use_reorderLinkControl = True , topo = topo )
114+
115+ # source routing table is passed to the plugin here
116+ endpointNIC .addPlugin ("sourceRoutingPlugin" , routing_table = _SR_routing_table )
117+ #### the following parameters will be passed down to linkcontrol via callback
118+ endpointNIC .link_bw = f"{ UNIFIED_ROUTER_LINK_BW }
119+ endpointNIC .input_buf_size = "32kB"
120+ endpointNIC .output_buf_size = "32kB"
121+ endpointNIC .vn_remap = [0 ]
122+
123+ # Set up VN remapping
124+ endpointNIC .vn_remap = [0 ]
125+
126+ ep = EmberMPIJob (0 ,topo .getNumNodes (), numCores = CPE * EPR )
127+ ep .setEndpointNIC (endpointNIC )
128+ ep .addMotif ("Init" )
129+ ep .addMotif (BENCH + BENCH_PARAMS )
130+ ep .addMotif ("Fini" )
131+ ep .nic .nic2host_lat = "10ns"
132+
133+ system = System ()
134+ system .setTopology (topo )
135+ system .allocateNodes (ep ,"linear" )
136+
137+ system .build ()
138+
139+ if __name__ == "__main__" :
140+ run_sst (* prepare_complete_4 ())
0 commit comments