[tuner] Convert compile_candidate.sh to py function compile_dispatch() #96
Conversation
There was a problem hiding this comment.
Seeing bash go away makes me happy!
I think this PR attempts to do two things at once:
- Reduce the amount of bash
- Make tuning less coupled to the exact model (Punet).
I think it would be easier to move forward if we focused on doing only one of these at a time. If we start with 2., we can define an interface for a model-specific tuners and still use bash scripts to implement that. Once that interface is in place, we can start removing bash.
The reason I think this should work out better is that the alternative is that we will come up with some short-lived python API for compilation/benchmarking that may not even survive the refactoring to make the tuner less coupled to punet.
The third alternative is to do both at the same time. This would also work, but I'm afraid it would be the most difficult to review.
| @@ -20,6 +20,8 @@ | |||
| import pickle | |||
| import iree.runtime as ireert | |||
| import random | |||
| import autotune_functions | |||
There was a problem hiding this comment.
How about we rename autotune.py to punet_autotune.py and then move the reusable code to an empty autotune.py?
| class CompilationMode(Enum): | ||
| DEFAULT = "default" | ||
| WINOGRAD = "winograd" |
There was a problem hiding this comment.
We don't need it anymore, punet doesn't use winograd
| prolog_content = prolog_file.read() | ||
| config_content = config_file.read() | ||
| epilog_content = epilog_file.read() | ||
|
|
||
| # Write the contents to the spec file | ||
| spec_file.write(prolog_content) | ||
| spec_file.write(config_content) | ||
| spec_file.write(epilog_content) |
There was a problem hiding this comment.
This is fine for now, but in the next iteration we should use the logic in generate_config.py instead.
| multiprocess_progress_wrapper( | ||
| num_worker=num_worker, task_list=task_list, function=run_command_wrapper | ||
| task_list = make_CompileDispatchTaskPack_task_list(args, path_config, candidates) | ||
| num_worker = max(min(args.max_cpu_workers, len(candidates)), 1) |
There was a problem hiding this comment.
| num_worker = max(min(args.max_cpu_workers, len(candidates)), 1) | |
| num_workers =min(args.max_cpu_workers, len(candidates)) |
If there are no candidates, we should exit
| assert input_path is not None | ||
| out_path = path_config.get_candidate_vmfb_path(candidate_id) | ||
|
|
||
| if result.success == False: |
There was a problem hiding this comment.
| if result.success == False: | |
| if not result.success: |
|
|
||
| if result.success == False: | ||
| failed_candidates.append(candidate_id) | ||
| # move candidates that failed to compile to a different dir. (i.e. candidates/compiled/#.mlir -> candidates/failed/#.mlir) |
There was a problem hiding this comment.
| # move candidates that failed to compile to a different dir. (i.e. candidates/compiled/#.mlir -> candidates/failed/#.mlir) | |
| # Move candidates that failed to compile to a different dir. (i.e. candidates/compiled/#.mlir -> candidates/failed/#.mlir) |
| out_path.unlink() | ||
| continue | ||
|
|
||
| # collect hash val for compiled files |
There was a problem hiding this comment.
| # collect hash val for compiled files | |
| # Collect hash values for all compiled files |
Bash scripts will be converted as functions in
autotune_functions.pyEach function will have a
TaskPackandResultPackas input and outputautotune_functions.py:
compile_dispatch()compiles a single dispatchautotune.py:
CompilationModeto replace the original list['default', 'winograd']compile_candidates()PathConfigCompileDispatchResultPackmaker function#_spec.mlirfor candidates