Skip to content

Commit

Permalink
address reviewer comments
Browse files Browse the repository at this point in the history 
Signed-off-by: Bangtian Liu <[email protected]>
  • Loading branch information
@bangtianliu
bangtianliu committed Jan 7, 2025
1 parent 4780a32 commit 8fe8476
Show file tree
Hide file tree
Showing 10 changed files with 234 additions and 225 deletions.
83 changes: 42 additions & 41 deletions tuner/examples/dispatch/dispatch_tuner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,48 +109,49 @@ def main():
path_config.base_dir.mkdir(parents=True, exist_ok=True)
path_config.output_unilog.touch()
candidate_trackers: list[libtuner.CandidateTracker] = []
mlir_ctx = ir.Context()
logger = logging.getLogger("tune")
tuner_context = TunerContext(mlir_ctx, logger)
dispatch_tuner = DispatchTuner(tuner_context)
stop_after_phase: str = args.stop_after

print("Setup logging")
libtuner.setup_logging(args, path_config)
print(path_config.run_log, end="\n\n")

if not args.dry_run:
print("Validating devices")
libtuner.validate_devices(args.devices)
print("Validation successful!\n")

print("Generating candidates...")
candidates = libtuner.generate_candidates(args, path_config, candidate_trackers)
print(f"Stored candidates in {path_config.candidates_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.generate_candidates:
return

print("Compiling candidates...")
compiled_candidates = libtuner.compile_dispatches(
args, path_config, candidates, candidate_trackers, dispatch_tuner
)
print(f"Compiled files are stored in {path_config.compiled_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.compile_dispatches:
return

print("Benchmarking compiled candidates...")
top_candidates = libtuner.benchmark_dispatches(
args, path_config, compiled_candidates, candidate_trackers, dispatch_tuner
)
print(f"\nStored results in {path_config.output_unilog.resolve()}\n")
if stop_after_phase == libtuner.ExecutionPhases.benchmark_dispatches:
return

libtuner.save_pickle(path_config.candidate_trackers_pkl, candidate_trackers)
print(f"Candidate trackers are saved in {path_config.candidate_trackers_pkl}\n")

print("Check the detailed execution logs in:")
print(path_config.run_log.resolve())
with TunerContext() as tuner_context:
dispatch_tuner = DispatchTuner(tuner_context)
stop_after_phase: str = args.stop_after

print("Setup logging")
libtuner.setup_logging(args, path_config)
print(path_config.run_log, end="\n\n")

if not args.dry_run:
print("Validating devices")
libtuner.validate_devices(args.devices)
print("Validation successful!\n")

print("Generating candidates...")
candidates = libtuner.generate_candidates(
args, path_config, candidate_trackers, tuner_context
)
print(f"Stored candidates in {path_config.candidates_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.generate_candidates:
return

print("Compiling candidates...")
compiled_candidates = libtuner.compile_dispatches(
args, path_config, candidates, candidate_trackers, dispatch_tuner
)
print(f"Compiled files are stored in {path_config.compiled_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.compile_dispatches:
return

print("Benchmarking compiled candidates...")
top_candidates = libtuner.benchmark_dispatches(
args, path_config, compiled_candidates, candidate_trackers, dispatch_tuner
)
print(f"\nStored results in {path_config.output_unilog.resolve()}\n")
if stop_after_phase == libtuner.ExecutionPhases.benchmark_dispatches:
return

libtuner.save_pickle(path_config.candidate_trackers_pkl, candidate_trackers)
print(f"Candidate trackers are saved in {path_config.candidate_trackers_pkl}\n")

print("Check the detailed execution logs in:")
print(path_config.run_log.resolve())

for candidate in candidate_trackers:
libtuner.logging.debug(candidate)
94 changes: 49 additions & 45 deletions tuner/examples/punet/punet_autotune.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

from tuner import libtuner
from pathlib import Path
from tuner.common import *


class PunetClient(libtuner.TuningClient):
Expand Down Expand Up @@ -142,51 +143,54 @@ def main():
print("Validation successful!\n")

print("Generating candidates...")
candidates = libtuner.generate_candidates(args, path_config, candidate_trackers)
print(f"Stored candidates in {path_config.candidates_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.generate_candidates:
return

print("Compiling candidates...")
compiled_candidates = libtuner.compile_dispatches(
args, path_config, candidates, candidate_trackers, punet_client
)
print(f"Compiled files are stored in {path_config.compiled_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.compile_dispatches:
return

print("Benchmarking compiled candidates...")
top_candidates = libtuner.benchmark_dispatches(
args, path_config, compiled_candidates, candidate_trackers, punet_client
)
print(f"Stored results in {path_config.output_unilog}\n")
if stop_after_phase == libtuner.ExecutionPhases.benchmark_dispatches:
return

print(f"Compiling top model candidates...")
punet_candidates = libtuner.compile_models(
args, path_config, top_candidates, candidate_trackers, punet_client
)
print(f"Model candidates compiled in {path_config.base_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.compile_models:
return

print("Benchmarking model candidates...")
libtuner.benchmark_models(
args, path_config, punet_candidates, candidate_trackers, punet_client
)
print(f"Stored results in {path_config.output_unilog}")
if stop_after_phase == libtuner.ExecutionPhases.benchmark_models:
return

libtuner.summerize_top_candidates(path_config, candidate_trackers)
print(f"Stored top candidates info in {path_config.result_summary_log}\n")

libtuner.save_pickle(path_config.candidate_trackers_pkl, candidate_trackers)
print(f"Candidate trackers are saved in {path_config.candidate_trackers_pkl}\n")

print("Check the detailed execution logs in:")
print(path_config.run_log)
with TunerContext() as tuner_context:
candidates = libtuner.generate_candidates(
args, path_config, candidate_trackers, tuner_context
)
print(f"Stored candidates in {path_config.candidates_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.generate_candidates:
return

print("Compiling candidates...")
compiled_candidates = libtuner.compile_dispatches(
args, path_config, candidates, candidate_trackers, punet_client
)
print(f"Compiled files are stored in {path_config.compiled_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.compile_dispatches:
return

print("Benchmarking compiled candidates...")
top_candidates = libtuner.benchmark_dispatches(
args, path_config, compiled_candidates, candidate_trackers, punet_client
)
print(f"Stored results in {path_config.output_unilog}\n")
if stop_after_phase == libtuner.ExecutionPhases.benchmark_dispatches:
return

print(f"Compiling top model candidates...")
punet_candidates = libtuner.compile_models(
args, path_config, top_candidates, candidate_trackers, punet_client
)
print(f"Model candidates compiled in {path_config.base_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.compile_models:
return

print("Benchmarking model candidates...")
libtuner.benchmark_models(
args, path_config, punet_candidates, candidate_trackers, punet_client
)
print(f"Stored results in {path_config.output_unilog}")
if stop_after_phase == libtuner.ExecutionPhases.benchmark_models:
return

libtuner.summerize_top_candidates(path_config, candidate_trackers)
print(f"Stored top candidates info in {path_config.result_summary_log}\n")

libtuner.save_pickle(path_config.candidate_trackers_pkl, candidate_trackers)
print(f"Candidate trackers are saved in {path_config.candidate_trackers_pkl}\n")

print("Check the detailed execution logs in:")
print(path_config.run_log)

for candidate in candidate_trackers:
libtuner.logging.debug(candidate)
Expand Down
116 changes: 57 additions & 59 deletions tuner/examples/test/tuner_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,65 +111,63 @@ def main():
print("Validation successful!\n")

print("Generating candidates...")
mlir_ctx = ir.Context()
logger = logging.getLogger("tune")
tuner_context = TunerContext(mlir_ctx, logger)
test_tuner = TestTuner(tuner_context)
candidates = libtuner.generate_candidate_specs(
args, path_config, candidate_trackers, test_tuner
)
print(f"Stored candidate specs in {path_config.specs_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.generate_candidates:
return

print("Compiling candidates...")
compiled_candidates = libtuner.compile(
args, path_config, candidates, candidate_trackers, test_tuner
)

print("Benchmarking compiled candidates...")
top_candidates = libtuner.benchmark(
args,
path_config,
compiled_candidates,
candidate_trackers,
test_tuner,
args.test_num_dispatch_candidates,
)

print("Compiling models with top candidates...")
test_tuner.compile_flags = [
"--iree-hal-target-backends=rocm",
f"--iree-hip-target={args.test_hip_target}",
]
compiled_model_candidates = libtuner.compile(
args,
path_config,
top_candidates,
candidate_trackers,
test_tuner,
args.test_model_file,
)

print("Benchmarking compiled model candidates...")
test_tuner.benchmark_flags = [
"--benchmark_repetitions=3",
"--input=2048x2048xf16",
"--input=2048x2048xf16",
]
top_model_candidates = libtuner.benchmark(
args,
path_config,
compiled_model_candidates,
candidate_trackers,
test_tuner,
args.test_num_model_candidates,
)

print(f"Top model candidates: {top_model_candidates}")

print("Check the detailed execution logs in:")
print(path_config.run_log.resolve())
with TunerContext() as tuner_context:
test_tuner = TestTuner(tuner_context)
candidates = libtuner.generate_candidate_specs(
args, path_config, candidate_trackers, test_tuner
)
print(f"Stored candidate specs in {path_config.specs_dir}\n")
if stop_after_phase == libtuner.ExecutionPhases.generate_candidates:
return

print("Compiling candidates...")
compiled_candidates = libtuner.compile(
args, path_config, candidates, candidate_trackers, test_tuner
)

print("Benchmarking compiled candidates...")
top_candidates = libtuner.benchmark(
args,
path_config,
compiled_candidates,
candidate_trackers,
test_tuner,
args.test_num_dispatch_candidates,
)

print("Compiling models with top candidates...")
test_tuner.compile_flags = [
"--iree-hal-target-backends=rocm",
f"--iree-hip-target={args.test_hip_target}",
]
compiled_model_candidates = libtuner.compile(
args,
path_config,
top_candidates,
candidate_trackers,
test_tuner,
args.test_model_file,
)

print("Benchmarking compiled model candidates...")
test_tuner.benchmark_flags = [
"--benchmark_repetitions=3",
"--input=2048x2048xf16",
"--input=2048x2048xf16",
]
top_model_candidates = libtuner.benchmark(
args,
path_config,
compiled_model_candidates,
candidate_trackers,
test_tuner,
args.test_num_model_candidates,
)

print(f"Top model candidates: {top_model_candidates}")

print("Check the detailed execution logs in:")
print(path_config.run_log.resolve())

for candidate in candidate_trackers:
libtuner.logging.debug(candidate)
Loading

0 comments on commit 8fe8476

Please sign in to comment.