add option to only benchmark schedule [run_process_replay] (#6204)

test/external/external_benchmark_schedule.py

@@ -10,6 +10,7 @@ if __name__ == "__main__":
 
   PROFILE = getenv("PROFILE", 1)
   FORWARD_ONLY = getenv("FORWARD_ONLY", 0)
+  SCHEDULE_ONLY = getenv("SCHEDULE_ONLY", 0)
 
   with Profiling(PROFILE):
     with Timing("***** model forward in "):
@@ -20,18 +21,19 @@ if __name__ == "__main__":
       with Timing("***** model schedule in "):
         sched = out.schedule()
 
-    asts = {x.ast.key:x.ast for x in sched if x.ast.op is UOps.SINK}.values()
-    kernels = []
-    with Profiling(PROFILE):
-      with Timing("***** model uops in "):
-        for ast in asts:
-          k = Kernel(ast)
-          k.hand_coded_optimizations()
-          kernels.append(k)
+    if not SCHEDULE_ONLY:
+      asts = {x.ast.key:x.ast for x in sched if x.ast.op is UOps.SINK}.values()
+      kernels = []
+      with Profiling(PROFILE):
+        with Timing("***** model uops in "):
+          for ast in asts:
+            k = Kernel(ast)
+            k.hand_coded_optimizations()
+            kernels.append(k)
 
-    with Profiling(PROFILE, fn="/tmp/schedule.prof"):
-      with Timing("***** model linearize in "):
-        for k in kernels: k.linearize()
+      with Profiling(PROFILE, fn="/tmp/schedule.prof"):
+        with Timing("***** model linearize in "):
+          for k in kernels: k.linearize()
 
     #renderer = Device[Device.DEFAULT].renderer
     #with Profiling(PROFILE, fn="/tmp/schedule.prof"):