fixes

python/ark/profiler.py

@@ -11,12 +11,21 @@
 
 def timeit(plan: Plan, iter: int, loop_mode: bool, warmup: int = 3):
     with Runtime() as rt:
-        rt.launch(plan=plan, loop_mode=loop_mode)
-        rt.run(iter=warmup)
-        start_time = time.time()
-        rt.run(iter=iter)
-        end_time = time.time()
-        return (end_time - start_time) / iter
+        if loop_mode:
+            rt.launch(plan=plan, loop_mode=loop_mode)
+            rt.run(iter=warmup)
+            rt.stop()
+            start_time = time.time()
+            rt.run(iter=iter)
+            elapsed = time.time() - start_time
+        else:
+            rt.launch(plan=plan, loop_mode=loop_mode)
+            rt.run(iter=warmup)
+            rt.stop()
+            rt.launch(plan=plan, loop_mode=loop_mode, record=True)
+            rt.run(iter=iter)
+            elapsed = rt.stop() / 1.0e3
+        return elapsed / iter
 
 
 class Profiler:

python/ark/runtime.py

@@ -65,6 +65,7 @@ def launch(
         device_id: int = -1,
         stream: int = 0,
         loop_mode: bool = True,
+        record: bool = False,
         tensor_mappings: Dict = {},
     ):
         """
@@ -93,7 +94,7 @@ def launch(
         exe = Executor.get()
         if plan_str != exe.plan() or device_id != exe.device_id():
             exe.compile(plan_str, device_id)
-        exe.launch(tensor_mappings, stream, loop_mode)
+        exe.launch(tensor_mappings, stream, loop_mode, record)
         self.state = Runtime.StateCode.LaunchedNotRunning
         self.loop_mode = loop_mode
 
@@ -148,7 +149,7 @@ def stop(self) -> float:
         """
         if not self.launched():
             log.WARN(f"ARK runtime is never launched, skip stopping")
-            return
+            return -1
         elapsed = Executor.get().stop()
         self.state = Runtime.StateCode.LaunchedNotRunning
         return elapsed