Skip to content

Commit

Permalink
pass existing tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@wmdi
wmdi committed Oct 16, 2024
1 parent da857a5 commit ef8c5c2
Show file tree
Hide file tree
Showing 4 changed files with 179 additions and 62 deletions.
18 changes: 10 additions & 8 deletions lib/compiler/test/src/compiler/machine_mapping/cost_estimator_for_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,32 +5,34 @@
namespace FlexFlow {

TestCostEstimator::TestCostEstimator(
std::function<float(OpCostEstimateKey const &)> const &get_operator_cost,
std::function<float(TensorSetMovement const &)> const
std::function<CostMetric(OpCostEstimateKey const &)> const
&get_operator_cost,
std::function<CostMetric(TensorSetMovement const &)> const
&get_communication_cost)
: get_operator_cost(get_operator_cost),
get_communication_cost(get_communication_cost) {}

float TestCostEstimator::estimate_cost(OpCostEstimateKey const &k) const {
CostMetric TestCostEstimator::estimate_cost(OpCostEstimateKey const &k) const {
return this->get_operator_cost(k);
}

float TestCostEstimator::estimate_cost(TensorSetMovement const &m) const {
CostMetric TestCostEstimator::estimate_cost(TensorSetMovement const &m) const {
return this->get_communication_cost(m);
}

CostEstimator make_fake_cost_estimator(
std::function<float(OpCostEstimateKey const &)> const &get_operator_cost,
std::function<float(TensorSetMovement const &)> const
std::function<CostMetric(OpCostEstimateKey const &)> const
&get_operator_cost,
std::function<CostMetric(TensorSetMovement const &)> const
&get_communication_cost) {

return CostEstimator::create<TestCostEstimator>(get_operator_cost,
get_communication_cost);
}

CostEstimator make_fake_cost_estimator(
std::unordered_map<OpCostEstimateKey, float> const &op_cost_map,
std::unordered_map<TensorSetMovement, float> const &comm_cost_map) {
std::unordered_map<OpCostEstimateKey, CostMetric> const &op_cost_map,
std::unordered_map<TensorSetMovement, CostMetric> const &comm_cost_map) {
return make_fake_cost_estimator(
[op_cost_map](OpCostEstimateKey const &k) { return op_cost_map.at(k); },
[comm_cost_map](TensorSetMovement const &m) {
Expand Down
18 changes: 9 additions & 9 deletions lib/compiler/test/src/compiler/machine_mapping/cost_estimator_for_test.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,27 +11,27 @@
namespace FlexFlow {

struct TestCostEstimator : public ICostEstimator {
std::function<float(OpCostEstimateKey const &)> get_operator_cost;
std::function<float(TensorSetMovement const &)> get_communication_cost;
std::function<CostMetric(OpCostEstimateKey const &)> get_operator_cost;
std::function<CostMetric(TensorSetMovement const &)> get_communication_cost;

TestCostEstimator() = delete;
TestCostEstimator(decltype(get_operator_cost) const &get_operator_cost,
decltype(get_communication_cost)
const &get_communication_cost);

float estimate_cost(OpCostEstimateKey const &) const override;

float estimate_cost(TensorSetMovement const &) const override;
CostMetric estimate_cost(OpCostEstimateKey const &) const override;
CostMetric estimate_cost(TensorSetMovement const &) const override;
};

CostEstimator make_fake_cost_estimator(
std::function<float(OpCostEstimateKey const &)> const &get_operator_cost,
std::function<float(TensorSetMovement const &)> const
std::function<CostMetric(OpCostEstimateKey const &)> const
&get_operator_cost,
std::function<CostMetric(TensorSetMovement const &)> const
&get_communication_cost);

CostEstimator make_fake_cost_estimator(
std::unordered_map<OpCostEstimateKey, float> const &op_cost_map,
std::unordered_map<TensorSetMovement, float> const &comm_cost_map);
std::unordered_map<OpCostEstimateKey, CostMetric> const &op_cost_map,
std::unordered_map<TensorSetMovement, CostMetric> const &comm_cost_map);

} // namespace FlexFlow

Expand Down
82 changes: 62 additions & 20 deletions lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,22 +118,22 @@ TEST_SUITE(FF_TEST_SUITE) {
}};

CostEstimator cost_estimator = make_fake_cost_estimator(
std::unordered_map<OpCostEstimateKey, float>{{
{map_unmapped_op_cost_estimate_key(k1, mv1), 1.0},
{map_unmapped_op_cost_estimate_key(k2, mv1), 2.0},
{map_unmapped_op_cost_estimate_key(k1, mv2), 1.5},
{map_unmapped_op_cost_estimate_key(k2, mv2), 2.5},
std::unordered_map<OpCostEstimateKey, CostMetric>{{
{map_unmapped_op_cost_estimate_key(k1, mv1), CostMetric{1.0, 1}},
{map_unmapped_op_cost_estimate_key(k2, mv1), CostMetric{2.0, 2}},
{map_unmapped_op_cost_estimate_key(k1, mv2), CostMetric{1.5, 3}},
{map_unmapped_op_cost_estimate_key(k2, mv2), CostMetric{2.5, 3}},
}},
std::unordered_map<TensorSetMovement, float>{{
{TensorSetMovement{{}}, 0.0},
std::unordered_map<TensorSetMovement, CostMetric>{{
{TensorSetMovement{{}}, CostMetric{0.0, 0}},
{concretize_abstracted_tensor_set_movement(movement1, mm1, mm1),
0.1},
CostMetric{0.1, 0}},
{concretize_abstracted_tensor_set_movement(movement1, mm2, mm2),
0.2},
CostMetric{0.2, 0}},
{concretize_abstracted_tensor_set_movement(movement1, mm1, mm2),
0.3},
CostMetric{0.3, 0}},
{concretize_abstracted_tensor_set_movement(movement1, mm2, mm1),
0.4},
CostMetric{0.4, 0}},
}});

MachineMappingContext context = MachineMappingContext{
Expand All @@ -150,11 +150,25 @@ TEST_SUITE(FF_TEST_SUITE) {
get_unconstrained_solution_for_layers(
get_all_leaf_paths(problem_tree));

MachineMappingResult result = get_optimal_machine_mapping(
cache, context, problem_tree, full_machine_spec, constraints);
MachineMemoryConstraints memory_constraints = MachineMemoryConstraints{
/*memory_limit=*/10,
};

MachineMappingConfig config = MachineMappingConfig{
/*enable_memory_optimization=*/false,
};

MachineMappingResult result =
get_optimal_machine_mapping(cache,
context,
problem_tree,
full_machine_spec,
constraints,
memory_constraints,
config);
MachineMappingResult correct = MachineMappingResult{
FeasibleMachineMappingResult{
/*runtime=*/1.0,
/*cost=*/CostMetric{1.0, 1},
/*machine_mapping=*/
ParallelLayerGuidObliviousMachineMapping{{
{binary_tree_root_path(), mv1},
Expand All @@ -173,11 +187,25 @@ TEST_SUITE(FF_TEST_SUITE) {
get_unconstrained_solution_for_layers(
get_all_leaf_paths(problem_tree));

MachineMappingResult result = get_optimal_machine_mapping(
cache, context, problem_tree, full_machine_spec, constraints);
MachineMemoryConstraints memory_constraints = MachineMemoryConstraints{
/*memory_limit=*/10,
};

MachineMappingConfig config = MachineMappingConfig{
/*enable_memory_optimization=*/false,
};

MachineMappingResult result =
get_optimal_machine_mapping(cache,
context,
problem_tree,
full_machine_spec,
constraints,
memory_constraints,
config);
MachineMappingResult correct = MachineMappingResult{
FeasibleMachineMappingResult{
/*runtime=*/1.0 + 2.0 + 0.1,
/*cost=*/CostMetric{1.0 + 2.0 + 0.1, 1 + 2 + 0},
/*machine_mapping=*/
ParallelLayerGuidObliviousMachineMapping{{
{
Expand Down Expand Up @@ -207,11 +235,25 @@ TEST_SUITE(FF_TEST_SUITE) {
get_unconstrained_solution_for_layers(
get_all_leaf_paths(problem_tree));

MachineMappingResult result = get_optimal_machine_mapping(
cache, context, problem_tree, full_machine_spec, constraints);
MachineMemoryConstraints memory_constraints = MachineMemoryConstraints{
/*memory_limit=*/10,
};

MachineMappingConfig config = MachineMappingConfig{
/*enable_memory_optimization=*/false,
};

MachineMappingResult result =
get_optimal_machine_mapping(cache,
context,
problem_tree,
full_machine_spec,
constraints,
memory_constraints,
config);
MachineMappingResult correct = MachineMappingResult{
FeasibleMachineMappingResult{
/*runtime=*/2.5,
/*cost=*/CostMetric{2.5, 3},
/*machine_mapping=*/
ParallelLayerGuidObliviousMachineMapping{{
{
Expand Down
Loading

0 comments on commit ef8c5c2

Please sign in to comment.