[Frontend] Change ExecutionEngine to PyTorchSimExecutionEngine

Scheduler/scheduler.py

@@ -140,7 +140,7 @@ def __str__(self):
     def register_model(model_name : str, compiled_model):
         SchedulerDNNModel.MODEL_MAP[model_name] = compiled_model
 
-class ExecutionEngine:
+class PyTorchSimRunner:
     PARTITION_BUSY = 0
     PARTITION_IDLE = 1
     SELECT_NOTHING = 2
@@ -262,7 +262,7 @@ def launch_kernel(self, current_cycle, partion_idx=0):
         self.partition_state[partion_idx] = self.PARTITION_BUSY
         return self.backend_simulator.launch(onnx_path, attribute_path, current_cycle, partion_idx)
 
-class FIFOExecutionEngine(ExecutionEngine):
+class FIFORunner(PyTorchSimRunner):
     def __init__(self, backend_simulator: BackendSimulator, num_partion=1) -> None:
         super().__init__(backend_simulator, num_partion)
 
@@ -297,7 +297,7 @@ def select_kernel(self, partition_idx):
         # No proper kernel now
         return self.SELECT_NOTHING
 
-class RRExecutionEngine(ExecutionEngine):
+class RoundRobinRunner(PyTorchSimRunner):
     def __init__(self, backend_simulator: BackendSimulator, num_partion=1) -> None:
         super().__init__(backend_simulator, num_partion)
         self.next_pointer = None
@@ -360,9 +360,9 @@ def __init__(self, num_request_queue=1, max_batch=1, engine_select=FIFO_ENGINE,
         self.backend_simulator = BackendSimulator(backend_path, backend_config)
         self.backend_simulator.interactive_simulation()
         if engine_select == Scheduler.FIFO_ENGINE:
-            self.execution_engine = FIFOExecutionEngine(self.backend_simulator, self.num_request_queue)
+            self.execution_engine = FIFORunner(self.backend_simulator, self.num_request_queue)
         elif engine_select == Scheduler.RR_ENGINE:
-            self.execution_engine = RRExecutionEngine(self.backend_simulator, self.num_request_queue)
+            self.execution_engine = RoundRobinRunner(self.backend_simulator, self.num_request_queue)
         else:
             print(f"Not supporetd engine type {engine_select}")
             exit(1)
@@ -480,7 +480,7 @@ def run(self, until_time):
         def execute_cycle():
             launch_ret_info = []
             for i in range(self.execution_engine.num_partion):
-                if self.execution_engine.partition_state[i] == ExecutionEngine.PARTITION_IDLE:
+                if self.execution_engine.partition_state[i] == PyTorchSimRunner.PARTITION_IDLE:
                     ret = self.execution_engine.launch_kernel(self.current_cycle, i)
                     launch_ret_info.append(ret)
 
@@ -495,7 +495,7 @@ def execute_cycle():
 
             for core_idx in result_list:
                 # Kernel is finished. So set idle state
-                self.execution_engine.partition_state[core_idx] = ExecutionEngine.PARTITION_IDLE
+                self.execution_engine.partition_state[core_idx] = PyTorchSimRunner.PARTITION_IDLE
 
             return result_list
 
@@ -543,4 +543,4 @@ def msec_to_cycle(self, msec):
             return msec
 
         freq = self.backend_simulator.get_core_freq()
-        return int(msec * (freq / 1000))
+        return int(msec * (freq / 1000))

experiments/gemm.py

@@ -48,7 +48,7 @@ def custom_matmul(a, b):
     if 'BACKENDSIM_SPIKE_ONLY' in os.environ:
         del os.environ['BACKENDSIM_SPIKE_ONLY']
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     run_matmul(size[0], size[1], size[2], config)

scripts/ILS_experiment/test_matmul.py

@@ -60,7 +60,7 @@ def custom_matmul(bias, a, b):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_matmul(device, *shape)

scripts/chiplet_prep.py

@@ -61,8 +61,8 @@ def modify_file(dump_path, name, address_numa_stride=None, subgraph_map=None):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     parser = argparse.ArgumentParser(description='Process folder argument.')
     parser.add_argument('size', type=int, help='Folder value', default=256)

test_extension_backend.py

@@ -22,8 +22,8 @@
 from tests.Fusion.test_matmul_activation import test_matmul_activation
 
 if __name__ == "__main__":
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     #test_vectoradd(device, (47, 10))
     #test_vector_scalar_add(device, (10, 10))

tests/Diffusion/test_diffusion.py

@@ -553,8 +553,8 @@ def test_upsample2d(
     args = parser.parse_args()
 
     sys.path.append(os.environ.get("TORCHSIM_DIR", "/workspace/PyTorchSim"))
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
 
     #test_upsample2d(device)

tests/Fusion/test_addmm_residual.py

@@ -43,8 +43,8 @@ def addmm_residual(a, b, c, d):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_addmm_residual(device, 32, 32, 32)
     test_addmm_residual(device, 128, 128, 128)

tests/Fusion/test_attention_fusion.py

@@ -75,8 +75,8 @@ def test_MHA(device, num_heads=12, embed_dim=768, input_seq=512):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_MHA(device)
     # test_Attention(device, head=16, seq=512, d_k=64)

tests/Fusion/test_bmm_reduction.py

@@ -42,8 +42,8 @@ def bmm(a, b):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     #test_bmm_reduce(device)
     test_bmm_reduce(device, 12, 512)

tests/Fusion/test_conv_fusion.py

@@ -101,8 +101,8 @@ def custom_conv_bn_relu(a, b, bias, c, d, e, f):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
 
     # Vanila test

tests/Fusion/test_matmul_activation.py

@@ -73,8 +73,8 @@ def test_matmul_activation(device, batch_size=16, input_size=32, output_size=8,
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_matmul_activation(device)
     test_matmul_activation(device, batch_size=32, input_size=32, output_size=32, activation_fn="sigmoid")

tests/Fusion/test_matmul_reduction.py

@@ -89,8 +89,8 @@ def matmul_fused(a, b, c, d):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_matmul_reduce(device, 3072, 512, 768)
     test_matmul_var_mean(device)

tests/Fusion/test_matmul_scalar.py

@@ -39,7 +39,7 @@ def matmul_fused(a, b, c):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_matmul_scalar(device)

tests/Fusion/test_prologue_fusion.py

@@ -88,8 +88,8 @@ def bmm(a, b, c, d):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_elem_broadcast_fusion(device)
     test_elem_fusion(device)

tests/Fusion/test_transformer_fusion.py

@@ -203,8 +203,8 @@ def test_EncoderBlock_validation(head=12, embed_dim=768, input_seq=512):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     #test_MHA(device)
     test_EncoderBlock(device)

tests/MLP/test_mlp.py

@@ -281,9 +281,9 @@ def train(model, device):
     return
 
 if __name__ == "__main__":
-    from Scheduler.scheduler import ExecutionEngine
+    from Scheduler.scheduler import PyTorchSimRunner
     torch.set_printoptions(threshold=float('inf'), linewidth=600)
-    module = ExecutionEngine.setup_device()
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
 
     test_mlp(device)

tests/MLP/test_mlp_cpu.py

@@ -399,7 +399,6 @@ def train(model, device):
 
 
 if __name__ == "__main__":
-    from Scheduler.scheduler import ExecutionEngine
     # torch.set_printoptions(threshold=float('inf'), linewidth=600)
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

tests/Mixtral_8x7B/test_attention.py

@@ -163,8 +163,8 @@ def test_rmsnorm(device, seq=32):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_rmsnorm(device, seq=1)
     test_concat(device, size1=(1, 8, 64, 64), size2=(1,8,1,64), dim=2)

tests/MoE/test_moe.py

@@ -783,9 +783,9 @@ def evaluation(model, evaluation_loader):
         train(opt_model, train_loader)
 
 if __name__ == "__main__":
-    from Scheduler.scheduler import ExecutionEngine
+    from Scheduler.scheduler import PyTorchSimRunner
     torch.set_printoptions(threshold=float('inf'), linewidth=600)
-    module = ExecutionEngine.setup_device()
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
 
     test_moe(device)

tests/test_activation.py

@@ -88,8 +88,8 @@ def test_SwiGLU(device, size=(128, 128)):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_ReLU(device, (47, 10))
     test_ReLU(device, (128, 128))

tests/test_add.py

@@ -58,8 +58,8 @@ def vectoradd(a, b):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_vectoradd(device, (1, 1))
     test_vectoradd(device, (47, 10))

tests/test_batchnorm.py

@@ -37,8 +37,8 @@ def test_BatchNorm(device, size=(1, 16, 64, 64)):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_BatchNorm(device)
     test_BatchNorm(device, size=(1,64, 32, 32))

tests/test_bmm.py

@@ -46,8 +46,8 @@ def bmm(a, b, bias):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_BMM(device)
     test_BMM(device, 2, 256, 128, 256)

tests/test_cnn.py

@@ -53,7 +53,7 @@ def test_CNN(device):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_CNN(device)

tests/test_conv2d.py

@@ -40,8 +40,8 @@ def custom_conv2d(a, b, bias):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     torch._dynamo.config.cache_size_limit = 64
     test_conv2d(device, batch_size=8, in_channels=3, out_channels=32, input_size=32, kernel_size=1, stride=1, padding=0)

tests/test_exponent.py

@@ -31,7 +31,7 @@ def exponent(a):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_exponent(device, size=(32, 32))

tests/test_indirect_access.py

@@ -48,8 +48,8 @@ def test_embedding(device, vocab_size, dim):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_indirect_vectoradd(device)
     #test_embedding(device, 1024, 2048)

tests/test_layernorm.py

@@ -41,8 +41,8 @@ def test_LayerNorm(device, size=(64, 64)):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     #test_LayerNorm(device)
     test_LayerNorm(device, shape)

tests/test_matmul.py

@@ -94,8 +94,8 @@ def custom_linear(a, b, bias):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_matmul(device, 32, 32, 32)
     test_matmul(device, 128, 128, 128)

tests/test_mlp.py

@@ -109,8 +109,8 @@ def test_optimizer(device):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_mlp(device)
     test_mlp_inf(device, batch_size=1, input_size=256, hidden_size=512, output_size=256)

tests/test_pool.py

@@ -47,8 +47,8 @@ def avgpool(a):
     import sys
     sys.path.append(os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim'))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     #test_maxpool(device, b=1, c=8, h=16, w=16)
     #test_maxpool(device, b=1, c=8, h=112, w=112)

tests/test_reduce.py

@@ -47,8 +47,8 @@ def reduce_sum(a, dim, keepdim):
     args = parser.parse_args()
     shape = tuple(map(int, args.shape.strip('()').split(',')))
 
-    from Scheduler.scheduler import ExecutionEngine
-    module = ExecutionEngine.setup_device()
+    from Scheduler.scheduler import PyTorchSimRunner
+    module = PyTorchSimRunner.setup_device()
     device = module.custom_device()
     test_reduce_sum(device, (29, 47), 1, keepdim=True)
     test_reduce_sum(device, (17, 68), 0, keepdim=True)