diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py index df7899e559..cec37b518b 100644 --- a/test/external/external_benchmark_schedule.py +++ b/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"):