fix bug (#79)

sarchlab · Jul 5, 2024 · 5d49f87 · 5d49f87
1 parent b26ab38
commit 5d49f87
Showing 1 changed file with 15 additions and 6 deletions.
diff --git a/samples/runner/report.go b/samples/runner/report.go
@@ -69,12 +69,21 @@ func (r *Runner) defineMetrics() {
 }
 
 func (r *Runner) addKernelTimeTracer() {
-	r.kernelTimeCounter = tracing.NewBusyTimeTracer(
-		r.platform.Engine,
-		func(task tracing.Task) bool {
-			return task.What == "*driver.LaunchKernelCommand"
-		})
-	tracing.CollectTrace(r.platform.Driver, r.kernelTimeCounter)
+	if *unifiedGPUFlag != "" {
+		r.kernelTimeCounter = tracing.NewBusyTimeTracer(
+			r.platform.Engine,
+			func(task tracing.Task) bool {
+				return task.What == "*driver.LaunchUnifiedMultiGPUKernelCommand"
+			})
+		tracing.CollectTrace(r.platform.Driver, r.kernelTimeCounter)
+	} else {
+		r.kernelTimeCounter = tracing.NewBusyTimeTracer(
+			r.platform.Engine,
+			func(task tracing.Task) bool {
+				return task.What == "*driver.LaunchKernelCommand"
+			})
+		tracing.CollectTrace(r.platform.Driver, r.kernelTimeCounter)
+	}
 
 	for _, gpu := range r.platform.GPUs {
 		gpuKernelTimeCounter := tracing.NewBusyTimeTracer(