[Scheduler] WIP (8)

Author: YWHyuk

PyTorchSimFrontend/extension_codecache.py

@@ -274,7 +274,7 @@ def run_kernel_simulation(*args, **kwargs):
                 # Dump arguments and meta data
                 dump_metadata(args, arg_attributes, result_path)
                 runtime_path = FunctionalSimulator.get_runtime_dump_path(result_path)
-                if extension_config.pytorchsim_functional_mode:
+                if extension_config.pytorchsim_functional_mode and not autotune:
                     funcsim = FunctionalSimulator(result_path, key)
                     funcsim.run_spike(args, arg_attributes,
                                     runtime_path, self.validation_binary_name,

PyTorchSimFrontend/mlir/mlir_autotune.py

@@ -61,12 +61,24 @@ def make_run_fn(
         # Check already cached result.
         write_path = get_write_path(self.source_code)
         key,  _ = write(self.source_code, "mlir", specified_dir=write_path)
-        result_path = os.path.join(extension_config.CONFIG_TORCHSIM_DUMP_PATH, "outputs", hash_prefix(key), "togsim_result/0")
-        if os.path.exists(result_path):
-            result = TOGSimulator.get_result_from_file(result_path)
-            def cached_run_fn(*args, **kwargs):
-                return result
-            return cached_run_fn
+        result_dir = os.path.join(extension_config.CONFIG_TORCHSIM_DUMP_PATH, "outputs", hash_prefix(key), "togsim_result")
+
+        # Find the most recent .log file in the result directory
+        if os.path.exists(result_dir) and os.path.isdir(result_dir):
+            log_files = [f for f in os.listdir(result_dir) if f.endswith('.log')]
+            if log_files:
+                # Sort by modification time, get the most recent file
+                log_files_with_time = [
+                    (f, os.path.getmtime(os.path.join(result_dir, f)))
+                    for f in log_files
+                ]
+                log_files_with_time.sort(key=lambda x: x[1], reverse=True)
+                latest_log_file = log_files_with_time[0][0]
+                result_path = os.path.join(result_dir, latest_log_file)
+                result = TOGSimulator.get_result_from_file(result_path)
+                def cached_run_fn(*args, **kwargs):
+                    return result
+                return cached_run_fn
 
         # Run a candidate code
         run_method = custom_async_compile.mlir(

Simulator/simulator.py

@@ -308,7 +308,7 @@ def _send_command(self, command_type, device_index, stream_index, tog_path="", a
             timestamp: Timestamp in nanoseconds (default: 0)
 
         Returns:
-            int: The kernel ID assigned to this command (or -1 for DEVICE_SYNC)
+            int: The kernel ID assigned to this command
         """
         if self.process is None:
             raise RuntimeError("[TOGSim] Simulator process is not running")

tests/test_scheduler.py

@@ -3,41 +3,25 @@
 import torch
 from torchvision.models import resnet18 as model1
 from test_transformer import EncoderBlock as model2
+from Simulator.simulator import TOGSimulator
 
 base_path = os.environ.get('TORCHSIM_DIR', default='/workspace/PyTorchSim')
-sys.path.append(base_path)
-from Scheduler.scheduler import Scheduler, SchedulerDNNModel, Request
 config = f'{base_path}/configs/systolic_ws_128x128_c2_simple_noc_tpuv3_partition.yml'
+os.environ['TOGSIM_CONFIG'] = config
 
 target_model1 = model1().eval()
 target_model2 = model2(768, 12).eval()
 
-# Init scheduler
-scheduler = Scheduler(num_request_queue=2, engine_select=Scheduler.FIFO_ENGINE, togsim_config=config)
-# Register compiled model
-opt_model1 = torch.compile(target_model1.to(device=scheduler.execution_engine.module.custom_device(), memory_format=torch.channels_last))
-opt_model2 = torch.compile(target_model2.to(device=scheduler.execution_engine.module.custom_device()))
-SchedulerDNNModel.register_model("resnet18", opt_model1)
-SchedulerDNNModel.register_model("bert", opt_model2)
-
-# Init input data
-model_input1 = torch.randn(1, 3, 224, 224)
-model_input2 = torch.randn(128, 768)
-
-# Init request
-new_request1 = Request("resnet18", [model_input1], [], request_queue_idx=0)
-new_request2 = Request("bert", [model_input2], [], request_queue_idx=1)
-new_request3 = Request("resnet18", [model_input1], [], request_queue_idx=0)
-new_request4 = Request("bert", [model_input2], [], request_queue_idx=1)
-
-# Add request to scheduler
-scheduler.add_request(new_request1, request_time=0)
-scheduler.add_request(new_request2, request_time=0)
-scheduler.add_request(new_request3, request_time=0)
-scheduler.add_request(new_request4, request_time=0)
-
-# Run scheduler
-while not scheduler.is_finished():
-    scheduler.schedule()
-
+device = torch.device("npu:0")
+opt_model1 = torch.compile(target_model1.to(device=device, memory_format=torch.channels_last))
+opt_model2 = torch.compile(target_model2.to(device=device))
+model_input1 = torch.randn(1, 3, 224, 224).to(device=device)
+model_input2 = torch.randn(128, 768).to(device=device)
+
+with TOGSimulator(config_path=config):
+    torch.npu.launch_model(opt_model1, model_input1, stream_index=0, timestamp=0)
+    torch.npu.launch_model(opt_model2, model_input2, stream_index=1, timestamp=0)
+    torch.npu.synchronize()
+    torch.npu.launch_model(opt_model1, model_input1, stream_index=0, timestamp=0)
+    torch.npu.launch_model(opt_model2, model_input2, stream_index=1, timestamp=0)
 print("Done")