Merge pull request bytedance#110 from bytedance/fix_micro_perf

[micro_perf] add barrier to avoid hang.
hliuca · Oct 7, 2024 · 94b142b · 94b142b
2 parents 26e28b1 + 13d8b00
commit 94b142b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/byte_micro_perf/backends/backend.py b/byte_micro_perf/backends/backend.py
@@ -187,15 +187,19 @@ def perf(self, input_shapes: List[List[int]], dtype):
                 prefer_iterations = self.iterations
 
                 # warmup
+                self.device_synchronize()
+                self.barrier()
                 for _ in range(warm_iterations):
                     self._run_operation(self.op, random.choice(tensor_list))
 
                 # test perf
                 self.device_synchronize()
+                self.barrier()
                 start_time = time.perf_counter_ns()
                 for i in range(test_iterations):
                     self._run_operation(self.op, random.choice(tensor_list))
                 self.device_synchronize()
+                self.barrier()
                 end_time = time.perf_counter_ns()
                 avg_op_duration = (end_time - start_time) / 1e9 / test_iterations
 
@@ -207,10 +211,12 @@ def perf(self, input_shapes: List[List[int]], dtype):
 
                 # perf
                 self.device_synchronize()
+                self.barrier()
                 start_time = time.perf_counter_ns()
                 for i in range(prefer_iterations):
                     self._run_operation(self.op, tensor_list[i % len(tensor_list)])
                 self.device_synchronize()
+                self.barrier()
                 end_time = time.perf_counter_ns()
 
                 # time in us