From a534303b37e673034bce7156e6593c2a00d9b30c Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Wed, 5 Aug 2026 11:43:29 -0700 Subject: [PATCH 1/5] Initial work on bin/compile-model. --- .proj.toml | 4 + CMakeLists.txt | 1 + bin/CMakeLists.txt | 4 + bin/compile-model/CMakeLists.txt | 10 ++ bin/compile-model/index.dox | 29 ++++++ bin/compile-model/src/compile-model/main.cc | 100 ++++++++++++++++++++ bin/index.dox | 2 + bin/run-model/index.dox | 29 ++++++ 8 files changed, 179 insertions(+) create mode 100644 bin/compile-model/CMakeLists.txt create mode 100644 bin/compile-model/index.dox create mode 100644 bin/compile-model/src/compile-model/main.cc create mode 100644 bin/run-model/index.dox diff --git a/.proj.toml b/.proj.toml index b40a92fea9..9411f57c13 100644 --- a/.proj.toml +++ b/.proj.toml @@ -100,6 +100,10 @@ has-cpu-only-benchmarks = false has-cuda-tests = false has-cuda-benchmarks = false +[targets.compile-model] +type = "bin" +cuda = false + [targets.export-model-arch] type = "bin" cuda = false diff --git a/CMakeLists.txt b/CMakeLists.txt index 4d62faf688..a0a154037a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -48,6 +48,7 @@ option(FF_BUILD_SUBSTITUTION_TOOL "build substitution conversion tool" OFF) option(FF_BUILD_VISUALIZATION_TOOL "build substitution visualization tool" ON) option(FF_BUILD_SP_IZATION_BENCHMARKING "build sp-ization benchmarking" ON) option(FF_BUILD_ARG_PARSER "build command line argument parser" OFF) +option(FF_BUILD_BIN_COMPILE_MODEL "build compile-model binary" ON) option(FF_BUILD_BIN_EXPORT_MODEL_ARCH "build export-model-arch utility" ON) option(FF_BUILD_BIN_RUN_MODEL "build run-model binary" ON) diff --git a/bin/CMakeLists.txt b/bin/CMakeLists.txt index 7bf7909320..f62b7d632c 100644 --- a/bin/CMakeLists.txt +++ b/bin/CMakeLists.txt @@ -10,6 +10,10 @@ if(FF_BUILD_SP_IZATION_BENCHMARKING) add_subdirectory(sp-ization-benchmarking) endif() +if(FF_BUILD_BIN_COMPILE_MODEL) + add_subdirectory(compile-model) +endif() + if(FF_BUILD_BIN_EXPORT_MODEL_ARCH) add_subdirectory(export-model-arch) endif() diff --git a/bin/compile-model/CMakeLists.txt b/bin/compile-model/CMakeLists.txt new file mode 100644 index 0000000000..769e7da4f9 --- /dev/null +++ b/bin/compile-model/CMakeLists.txt @@ -0,0 +1,10 @@ +ff_add_executable( + NAME + compile-model + SRC_PATTERNS + src/*.cc + PRIVATE_INCLUDE + include/ + DEPS + compiler +) diff --git a/bin/compile-model/index.dox b/bin/compile-model/index.dox new file mode 100644 index 0000000000..e93548fc4b --- /dev/null +++ b/bin/compile-model/index.dox @@ -0,0 +1,29 @@ +namespace FlexFlow { +/** + +@page compile-model compile-model + +\brief Run the model specified by the mapped PCG, encoded as JSON. + +The main entry point for the %FlexFlow %Train compiler. The compiler takes a computation graph describing the operations to be performed, an transforms it into a mapped parallel computation graph (PCG) that describes how to do so on a parallel machine. + +\verbatim +$ proj cmake # if you haven't already +... +$ proj build +... +$ ./build/normal/bin/compile-model/compile-model -h +\endverbatim + +The above should print the help message for `compile-model`. + +A typical workflow would involve generating a model (e.g., via \ref export-model-arch), compiling it, and then running (via \ref run-model): + +\verbatim +$ ./build/normal/bin/export-model-arch/export-model-arch split_test > split_test_cg.json +$ ./build/normal/bin/compile-model/compile-model split_test_cg.json > split_test_mpcg.json +$ ./build/normal/bin/run-model-arch/run-model-arch split_test_mpcg.json +\endverbatim + +*/ +} diff --git a/bin/compile-model/src/compile-model/main.cc b/bin/compile-model/src/compile-model/main.cc new file mode 100644 index 0000000000..0abc657ebe --- /dev/null +++ b/bin/compile-model/src/compile-model/main.cc @@ -0,0 +1,100 @@ +#include "pcg/file_format/v1/v1_computation_graph.h" +#include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.h" +#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.h" +#include "pcg/pcg_from_computation_graph.h" +#include "utils/cli/cli_get_help_message.h" +#include "utils/cli/cli_parse.h" +#include "utils/cli/cli_parse_result.h" +#include "utils/cli/cli_spec.h" +#include + +using namespace FlexFlow; + +int main(int argc, char **argv) { + CLISpec cli = empty_cli_spec(); + + CLIArgumentKey arg_key_help = cli_add_help_flag(cli); + + CLIArgumentKey key_cg_json = cli_add_positional_argument( + cli, + CLIPositionalArgumentSpec{ + "cg_json", + std::nullopt, + "path to a file containing computation graph, encoded as JSON"}); + + CLIArgumentKey key_mpcg_json_output = cli_add_positional_argument( + cli, + CLIPositionalArgumentSpec{ + "mpcg_json_output", + std::nullopt, + "path to write the resulting mapping PCG, encoded as JSON"}); + + std::vector strategy_options = {"passthrough", + "data_parallel", + "unity", + "mcmc"}; + CLIArgumentKey key_strategy = cli_add_positional_argument( + cli, + CLIPositionalArgumentSpec{ + "strategy", strategy_options, "compilation strategy for building the mapped PCG"}); + + ASSERT(argc >= 1); + std::string prog_name = argv[0]; + + CLIParseResult parsed = ({ + tl::expected result = + cli_parse(cli, argc, argv); + if (!result.has_value()) { + std::string error_msg = result.error(); + std::cerr << cli_get_help_message(prog_name, cli); + std::cerr << std::endl; + std::cerr << "error: " << error_msg << std::endl; + return 1; + } + + result.value(); + }); + + bool help = cli_get_flag(parsed, arg_key_help); + if (help) { + std::cerr << cli_get_help_message(prog_name, cli); + return 1; + } + + std::string cg_json = cli_get_argument(parsed, key_cg_json); + std::string mpcg_json_output = cli_get_argument(parsed, key_mpcg_json_output); + std::string strategy = cli_get_argument(parsed, key_strategy); + + ComputationGraph cg = [&]() { + std::ifstream f{cg_json}; + nlohmann::json cg_json = nlohmann::json::parse(f); + return from_v1(cg_json.get()); + }(); + + MappedParallelComputationGraph mpcg = [&]() { + if (strategy == "passthrough") { + ParallelComputationGraph pcg = pcg_from_computation_graph(cg); + std::map + mapped_op_task_groups; + return mapped_pcg_from_pcg_and_mapped_op_task_groups( + pcg, + mapped_op_task_groups); + } else if (strategy == "data_parallel") { + NOT_IMPLEMENTED(); + } else if (strategy == "unity") { + NOT_IMPLEMENTED(); + } else if (strategy == "mcmc") { + NOT_IMPLEMENTED(); + } else { + PANIC("no such strategy: {}", strategy); + } + }(); + + { + std::ofstream f{mpcg_json_output}; + nlohmann::json mpcg_json = to_v1(mpcg); + f << mpcg_json; + } + + return 0; +} diff --git a/bin/index.dox b/bin/index.dox index 4944e50067..e4d8455d6d 100644 --- a/bin/index.dox +++ b/bin/index.dox @@ -4,8 +4,10 @@ This directory contains command-line interfaces for %FlexFlow %Train and associated tools (all in C++). +- \subpage compile-model "": \copybrief compile-model - \subpage export-model-arch "": \copybrief export-model-arch - \subpage protobuf-to-json "": \copybrief protobuf-to-json +- \subpage run-model "": \copybrief run-model - \subpage sp-ization-benchmarking "": \copybrief sp-ization-benchmarking - \subpage substitution-to-dot "": \copybrief substitution-to-dot diff --git a/bin/run-model/index.dox b/bin/run-model/index.dox new file mode 100644 index 0000000000..90154551a7 --- /dev/null +++ b/bin/run-model/index.dox @@ -0,0 +1,29 @@ +namespace FlexFlow { +/** + +@page run-model run-model + +\brief Run the model specified by the mapped PCG, encoded as JSON. + +The main entry point for the %FlexFlow %Train model runner. The runner takes a mapped PCG generated by the compiler and runs it on a (potentially distributed, potentially heterogeneous) machine via the Realm runtime layer. + +\verbatim +$ proj cmake # if you haven't already +... +$ proj build +... +$ ./build/normal/bin/run-model/run-model -h +\endverbatim + +The above should print the help message for `run-model`. + +A typical workflow would involve generating a model (e.g., via \ref export-model-arch) and then compiling it (via \ref compile-model) to a mapped PCG before running: + +\verbatim +$ ./build/normal/bin/export-model-arch/export-model-arch split_test > split_test_cg.json +$ ./build/normal/bin/compile-model/compile-model split_test_cg.json > split_test_mpcg.json +$ ./build/normal/bin/run-model-arch/run-model-arch split_test_mpcg.json +\endverbatim + +*/ +} From 2d8c72592ae41bd2dac1c7226f54776d3dc7b468 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 6 Aug 2026 11:05:58 -0700 Subject: [PATCH 2/5] Format. --- bin/compile-model/src/compile-model/main.cc | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/bin/compile-model/src/compile-model/main.cc b/bin/compile-model/src/compile-model/main.cc index 0abc657ebe..c52d1d7971 100644 --- a/bin/compile-model/src/compile-model/main.cc +++ b/bin/compile-model/src/compile-model/main.cc @@ -29,14 +29,14 @@ int main(int argc, char **argv) { std::nullopt, "path to write the resulting mapping PCG, encoded as JSON"}); - std::vector strategy_options = {"passthrough", - "data_parallel", - "unity", - "mcmc"}; + std::vector strategy_options = { + "passthrough", "data_parallel", "unity", "mcmc"}; CLIArgumentKey key_strategy = cli_add_positional_argument( cli, CLIPositionalArgumentSpec{ - "strategy", strategy_options, "compilation strategy for building the mapped PCG"}); + "strategy", + strategy_options, + "compilation strategy for building the mapped PCG"}); ASSERT(argc >= 1); std::string prog_name = argv[0]; @@ -77,8 +77,7 @@ int main(int argc, char **argv) { std::map mapped_op_task_groups; return mapped_pcg_from_pcg_and_mapped_op_task_groups( - pcg, - mapped_op_task_groups); + pcg, mapped_op_task_groups); } else if (strategy == "data_parallel") { NOT_IMPLEMENTED(); } else if (strategy == "unity") { From 67f2441337c4815dff54e32ebe1aa0dddbb9f1e6 Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 6 Aug 2026 11:58:40 -0700 Subject: [PATCH 3/5] Finish passthrough with mapping to one device. --- bin/compile-model/index.dox | 4 +-- bin/compile-model/src/compile-model/main.cc | 28 ++++++++++++++++++++- bin/run-model/index.dox | 4 +-- 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/bin/compile-model/index.dox b/bin/compile-model/index.dox index e93548fc4b..8255278d35 100644 --- a/bin/compile-model/index.dox +++ b/bin/compile-model/index.dox @@ -21,8 +21,8 @@ A typical workflow would involve generating a model (e.g., via \ref export-model \verbatim $ ./build/normal/bin/export-model-arch/export-model-arch split_test > split_test_cg.json -$ ./build/normal/bin/compile-model/compile-model split_test_cg.json > split_test_mpcg.json -$ ./build/normal/bin/run-model-arch/run-model-arch split_test_mpcg.json +$ ./build/normal/bin/compile-model/compile-model split_test_cg.json split_test_mpcg.json passthrough +$ ./build/normal/bin/run-model/run-model split_test_mpcg.json \endverbatim */ diff --git a/bin/compile-model/src/compile-model/main.cc b/bin/compile-model/src/compile-model/main.cc index c52d1d7971..2dc7328839 100644 --- a/bin/compile-model/src/compile-model/main.cc +++ b/bin/compile-model/src/compile-model/main.cc @@ -6,6 +6,10 @@ #include "utils/cli/cli_parse.h" #include "utils/cli/cli_parse_result.h" #include "utils/cli/cli_spec.h" +#include "utils/containers/binary_merge_disjoint_maps.h" +#include "utils/containers/map_from_pairs.h" +#include "utils/containers/map_values.h" +#include "utils/containers/transform.h" #include using namespace FlexFlow; @@ -73,9 +77,31 @@ int main(int argc, char **argv) { MappedParallelComputationGraph mpcg = [&]() { if (strategy == "passthrough") { + auto mapping_for_pcg_invocation_info = + [](ParallelLayerInvocationInfo const &info) { + // Everything maps to zero + ParallelTensorSpaceCoordinate tensor_coord_zero{ + 0_n, 0_n, FFOrdered{0_n}}; + MachineSpaceCoordinate machine_coord_zero{0_n, 0_n}; + OperatorAtomicTaskShardBinding shard_binding{ + binary_merge_disjoint_maps( + map_values(info.incoming, + [&](ParallelTensorInfo const &) { + return tensor_coord_zero; + }), + map_values(info.outgoing, [&](ParallelTensorInfo const &) { + return tensor_coord_zero; + }))}; + return std::pair{ + info.layer_info.guid, + MappedOperatorTaskGroup{{{machine_coord_zero, shard_binding}}}, + }; + }; ParallelComputationGraph pcg = pcg_from_computation_graph(cg); std::map - mapped_op_task_groups; + mapped_op_task_groups = + map_from_pairs(transform(pcg_get_invocation_info_set(pcg), + mapping_for_pcg_invocation_info)); return mapped_pcg_from_pcg_and_mapped_op_task_groups( pcg, mapped_op_task_groups); } else if (strategy == "data_parallel") { diff --git a/bin/run-model/index.dox b/bin/run-model/index.dox index 90154551a7..e5ae36db66 100644 --- a/bin/run-model/index.dox +++ b/bin/run-model/index.dox @@ -21,8 +21,8 @@ A typical workflow would involve generating a model (e.g., via \ref export-model \verbatim $ ./build/normal/bin/export-model-arch/export-model-arch split_test > split_test_cg.json -$ ./build/normal/bin/compile-model/compile-model split_test_cg.json > split_test_mpcg.json -$ ./build/normal/bin/run-model-arch/run-model-arch split_test_mpcg.json +$ ./build/normal/bin/compile-model/compile-model split_test_cg.json split_test_mpcg.json passthrough +$ ./build/normal/bin/run-model/run-model split_test_mpcg.json \endverbatim */ From 9d630445a63959b1ef81bfe8db407ccf27686f2c Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 6 Aug 2026 13:49:54 -0700 Subject: [PATCH 4/5] Refactor to clarify code. --- bin/compile-model/src/compile-model/main.cc | 58 +++++++++++---------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/bin/compile-model/src/compile-model/main.cc b/bin/compile-model/src/compile-model/main.cc index 2dc7328839..9fe45b90bf 100644 --- a/bin/compile-model/src/compile-model/main.cc +++ b/bin/compile-model/src/compile-model/main.cc @@ -14,6 +14,36 @@ using namespace FlexFlow; +static std::pair + single_device_mapping_from_pcg_invocation_info( + ParallelLayerInvocationInfo const &info) { + // Everything maps to zero + ParallelTensorSpaceCoordinate tensor_coord_zero{ + 0_n, 0_n, FFOrdered{0_n}}; + MachineSpaceCoordinate machine_coord_zero{0_n, 0_n}; + OperatorAtomicTaskShardBinding shard_binding{binary_merge_disjoint_maps( + map_values(info.incoming, + [&](ParallelTensorInfo const &) { return tensor_coord_zero; }), + map_values(info.outgoing, [&](ParallelTensorInfo const &) { + return tensor_coord_zero; + }))}; + return std::pair{ + info.layer_info.guid, + MappedOperatorTaskGroup{{{machine_coord_zero, shard_binding}}}, + }; +} + +static MappedParallelComputationGraph + lift_cg_to_mpcg_for_single_device(ComputationGraph const &cg) { + ParallelComputationGraph pcg = pcg_from_computation_graph(cg); + std::map + mapped_op_task_groups = map_from_pairs( + transform(pcg_get_invocation_info_set(pcg), + single_device_mapping_from_pcg_invocation_info)); + return mapped_pcg_from_pcg_and_mapped_op_task_groups(pcg, + mapped_op_task_groups); +} + int main(int argc, char **argv) { CLISpec cli = empty_cli_spec(); @@ -77,33 +107,7 @@ int main(int argc, char **argv) { MappedParallelComputationGraph mpcg = [&]() { if (strategy == "passthrough") { - auto mapping_for_pcg_invocation_info = - [](ParallelLayerInvocationInfo const &info) { - // Everything maps to zero - ParallelTensorSpaceCoordinate tensor_coord_zero{ - 0_n, 0_n, FFOrdered{0_n}}; - MachineSpaceCoordinate machine_coord_zero{0_n, 0_n}; - OperatorAtomicTaskShardBinding shard_binding{ - binary_merge_disjoint_maps( - map_values(info.incoming, - [&](ParallelTensorInfo const &) { - return tensor_coord_zero; - }), - map_values(info.outgoing, [&](ParallelTensorInfo const &) { - return tensor_coord_zero; - }))}; - return std::pair{ - info.layer_info.guid, - MappedOperatorTaskGroup{{{machine_coord_zero, shard_binding}}}, - }; - }; - ParallelComputationGraph pcg = pcg_from_computation_graph(cg); - std::map - mapped_op_task_groups = - map_from_pairs(transform(pcg_get_invocation_info_set(pcg), - mapping_for_pcg_invocation_info)); - return mapped_pcg_from_pcg_and_mapped_op_task_groups( - pcg, mapped_op_task_groups); + return lift_cg_to_mpcg_for_single_device(cg); } else if (strategy == "data_parallel") { NOT_IMPLEMENTED(); } else if (strategy == "unity") { From 3c53e28a94217c998ce252ab5ad89b46607dea1e Mon Sep 17 00:00:00 2001 From: Elliott Slaughter Date: Thu, 6 Aug 2026 16:36:35 -0700 Subject: [PATCH 5/5] Implement the compiler-based paths. --- bin/compile-model/CMakeLists.txt | 2 + bin/compile-model/src/compile-model/main.cc | 125 ++++++++++++++++++-- 2 files changed, 120 insertions(+), 7 deletions(-) diff --git a/bin/compile-model/CMakeLists.txt b/bin/compile-model/CMakeLists.txt index 769e7da4f9..5d7a7e0212 100644 --- a/bin/compile-model/CMakeLists.txt +++ b/bin/compile-model/CMakeLists.txt @@ -7,4 +7,6 @@ ff_add_executable( include/ DEPS compiler + kernels + local-execution ) diff --git a/bin/compile-model/src/compile-model/main.cc b/bin/compile-model/src/compile-model/main.cc index 9fe45b90bf..75e69238af 100644 --- a/bin/compile-model/src/compile-model/main.cc +++ b/bin/compile-model/src/compile-model/main.cc @@ -1,3 +1,14 @@ +#include "compiler/compiler.h" +#include "compiler/cost_estimator/cost_estimator.h" +#include "compiler/data_parallelism/data_parallelism_config.dtg.h" +#include "compiler/mcmc/mcmc_over_mapped_pcg_config.dtg.h" +#include "compiler/search_result.h" +#include "compiler/unity_algorithm/unity_search_config.dtg.h" +#include "kernels/allocation.h" +#include "kernels/device_handle_t.h" +#include "kernels/local_cpu_allocator.h" +#include "kernels/local_cuda_allocator.h" +#include "local-execution/cost_estimator/local_cost_estimator.h" #include "pcg/file_format/v1/v1_computation_graph.h" #include "pcg/file_format/v1/v1_mapped_parallel_computation_graph.h" #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.h" @@ -10,7 +21,10 @@ #include "utils/containers/map_from_pairs.h" #include "utils/containers/map_values.h" #include "utils/containers/transform.h" +#include "utils/optional.h" +#include "utils/positive_int/positive_int.h" #include +#include using namespace FlexFlow; @@ -44,6 +58,94 @@ static MappedParallelComputationGraph mapped_op_task_groups); } +static MachineSpecification discover_machine() { + // TODO (Elliott): actually discover the machine topology + MachineComputeSpecification compute{/*num_nodes=*/1_p, + /*num_cpus_per_node=*/2_p, + /*num_gpus_per_node=*/2_p}; + MachineInterconnectSpecification interconnect{ + /*inter_node_bandwidth=*/bytes_per_second_t{1.0}, + /*intra_node_bandwidth=*/bytes_per_second_t{2.0}}; + MachineSpecification machine{compute, interconnect}; + return machine; +} + +static Allocator create_allocator(bool cpu_only) { + if (cpu_only) { + return create_local_cpu_memory_allocator(); + } else { + return create_local_cuda_memory_allocator(); + } +} + +static std::optional + create_device_handle(bool cpu_only) { + if (cpu_only) { + return std::nullopt; + } + + return initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); +} + +static device_handle_t create_device_handle( + bool cpu_only, + std::optional const &managed_handle) { + if (cpu_only) { + return cpu_make_device_handle_t(); + } else { + return gpu_make_device_handle_t(assert_unwrap(managed_handle).raw_handle()); + } +} + +static CostEstimator create_cost_estimator( + MachineSpecification const &machine, + bool cpu_only, + std::optional const &managed_handle) { + Allocator allocator = create_allocator(cpu_only); + device_handle_t ff_handle = create_device_handle(cpu_only, managed_handle); + global_device_id_t global_device_id = global_device_id_t{ + /*coord=*/MachineSpaceCoordinate{ + /*node_idx=*/0_n, + /*device_idx=*/0_n, + }, + /*device_type=*/(cpu_only ? DeviceType::CPU : DeviceType::GPU), + }; + return CostEstimator::create( + machine.interconnect_specification, + allocator, + ProfilingSettings{/*warmup_iters=*/2, /*measure_iters=*/5}, + ff_handle, + global_device_id); +} + +static AlgorithmConfig + select_compiler_algorithm(std::string const &strategy, + bool cpu_only, + MachineSpecification const &machine) { + if (strategy == "data_parallel") { + positive_int degree = + machine.compute_specification.num_nodes * + (cpu_only ? machine.compute_specification.num_cpus_per_node + : machine.compute_specification.num_gpus_per_node); + return AlgorithmConfig{ + DataParallelismConfig{/*degree=*/degree.int_from_positive_int()}}; + } else if (strategy == "unity") { + // TODO: pick better defaults + return AlgorithmConfig{ + UnitySearchConfig{/*alpha=*/0.5, /*budget=*/100, /*max_num_ops=*/100}}; + } else if (strategy == "mcmc") { + // TODO: pick better defaults + return AlgorithmConfig{ + MCMCOverMappedPCGConfig{/*temperature=*/0.5, + /*num_iterations=*/100_n, + /*substitution_frequency=*/0.5}}; + } else { + PANIC("no such strategy: {}", strategy); + } +} + int main(int argc, char **argv) { CLISpec cli = empty_cli_spec(); @@ -72,6 +174,11 @@ int main(int argc, char **argv) { strategy_options, "compilation strategy for building the mapped PCG"}); + CLIArgumentKey key_cpu = cli_add_flag( + cli, + CLIFlagSpec{ + "cpu", std::nullopt, "optimize graph for CPUs only (no GPUs)"}); + ASSERT(argc >= 1); std::string prog_name = argv[0]; @@ -98,6 +205,7 @@ int main(int argc, char **argv) { std::string cg_json = cli_get_argument(parsed, key_cg_json); std::string mpcg_json_output = cli_get_argument(parsed, key_mpcg_json_output); std::string strategy = cli_get_argument(parsed, key_strategy); + bool cpu = cli_get_flag(parsed, key_cpu); ComputationGraph cg = [&]() { std::ifstream f{cg_json}; @@ -108,14 +216,17 @@ int main(int argc, char **argv) { MappedParallelComputationGraph mpcg = [&]() { if (strategy == "passthrough") { return lift_cg_to_mpcg_for_single_device(cg); - } else if (strategy == "data_parallel") { - NOT_IMPLEMENTED(); - } else if (strategy == "unity") { - NOT_IMPLEMENTED(); - } else if (strategy == "mcmc") { - NOT_IMPLEMENTED(); } else { - PANIC("no such strategy: {}", strategy); + MachineSpecification machine = discover_machine(); + // Need to root this on the stack so it stays alive for the whole session + std::optional managed_handle = + create_device_handle(cpu); + CostEstimator estimator = + create_cost_estimator(machine, cpu, managed_handle); + AlgorithmConfig algorithm = + select_compiler_algorithm(strategy, cpu, machine); + SearchResult result = optimize(cg, machine, estimator, algorithm); + return get_mapped_pcg_from_search_result(result); } }();