Skip to content

Change OpCostMetrics.memory to be a nonnegative_int #1586

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lockshaw merged 3 commits into from
Jan 21, 2025
Merged

Change OpCostMetrics.memory to be a nonnegative_int #1586

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
392d8bc
Change OpCostMetrics.memory to be a nonnegative_int
Jan 21, 2025
e8e9e5f
Merge branch 'master' into issue-1584
victorli2002 Jan 21, 2025
efc73f5
Change OpCostMetrics.memory to be a nonnegative_int
Jan 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion lib/compiler/include/compiler/cost_estimator/op_cost_metrics.struct.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ features = [
]

includes = [
"utils/nonnegative_int/nonnegative_int.h"
]

[[fields]]
Expand All @@ -15,4 +16,4 @@ type = "float"

[[fields]]
name = "memory"
type = "size_t"
type = "::FlexFlow::nonnegative_int"
8 changes: 4 additions & 4 deletions lib/compiler/test/src/compiler/machine_mapping/get_optimal_machine_mapping.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,13 +146,13 @@ TEST_SUITE(FF_TEST_SUITE) {

auto map1 = std::unordered_map<OpCostEstimateKey, OpCostMetrics>{{
{map_unmapped_op_cost_estimate_key(k1, mv1),
OpCostMetrics{/*runtime=*/1.0, /*memory=*/0}},
OpCostMetrics{/*runtime=*/1.0, /*memory=*/nonnegative_int{0}}},
{map_unmapped_op_cost_estimate_key(k2, mv1),
OpCostMetrics{/*runtime=*/2.0, /*memory=*/0}},
OpCostMetrics{/*runtime=*/2.0, /*memory=*/nonnegative_int{0}}},
{map_unmapped_op_cost_estimate_key(k1, mv2),
OpCostMetrics{/*runtime=*/1.5, /*memory=*/0}},
OpCostMetrics{/*runtime=*/1.5, /*memory=*/nonnegative_int{0}}},
{map_unmapped_op_cost_estimate_key(k2, mv2),
OpCostMetrics{/*runtime=*/2.5, /*memory=*/0}},
OpCostMetrics{/*runtime=*/2.5, /*memory=*/nonnegative_int{0}}},
}};

CostEstimator cost_estimator = make_fake_cost_estimator(
Expand Down
22 changes: 13 additions & 9 deletions ...c/compiler/machine_mapping/memory_optimization/get_optimal_machine_mapping_with_memory.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,10 +146,14 @@ TEST_SUITE(FF_TEST_SUITE) {

CostEstimator cost_estimator = make_fake_cost_estimator(
std::unordered_map<OpCostEstimateKey, OpCostMetrics>{{
{map_unmapped_op_cost_estimate_key(k1, mv1), OpCostMetrics{1.0, 2}},
{map_unmapped_op_cost_estimate_key(k2, mv1), OpCostMetrics{2.0, 3}},
{map_unmapped_op_cost_estimate_key(k1, mv2), OpCostMetrics{1.5, 1}},
{map_unmapped_op_cost_estimate_key(k2, mv2), OpCostMetrics{2.5, 2}},
{map_unmapped_op_cost_estimate_key(k1, mv1),
OpCostMetrics{1.0, nonnegative_int{2}}},
{map_unmapped_op_cost_estimate_key(k2, mv1),
OpCostMetrics{2.0, nonnegative_int{3}}},
{map_unmapped_op_cost_estimate_key(k1, mv2),
OpCostMetrics{1.5, nonnegative_int{1}}},
{map_unmapped_op_cost_estimate_key(k2, mv2),
OpCostMetrics{2.5, nonnegative_int{2}}},
}},
std::unordered_map<TensorSetMovement, float>{{
{TensorSetMovement{{}}, 0.0},
Expand Down Expand Up @@ -183,13 +187,13 @@ TEST_SUITE(FF_TEST_SUITE) {
cache, context, problem_tree, full_machine_spec, constraints);
MachineMappingWithMemoryResult correct = MachineMappingWithMemoryResult{{
MachineMappingForSingleLayer{
OpCostMetrics{1.0, 2},
OpCostMetrics{1.0, nonnegative_int{2}},
ParallelLayerGuidObliviousMachineMapping{{
{binary_tree_root_path(), mv1},
}},
},
MachineMappingForSingleLayer{
OpCostMetrics{1.5, 1},
OpCostMetrics{1.5, nonnegative_int{1}},
ParallelLayerGuidObliviousMachineMapping{{
{binary_tree_root_path(), mv2},
}},
Expand All @@ -214,7 +218,7 @@ TEST_SUITE(FF_TEST_SUITE) {
MachineMappingForSingleLayer{
OpCostMetrics{
/*runtime=*/1.0 + 2.0 + 0.1,
/*memory=*/2 + 3,
/*memory=*/nonnegative_int{2 + 3},
},
ParallelLayerGuidObliviousMachineMapping{{
{
Expand All @@ -232,7 +236,7 @@ TEST_SUITE(FF_TEST_SUITE) {
}},
},
MachineMappingForSingleLayer{
OpCostMetrics{1.5 + 2.5 + 0.1, 1 + 2},
OpCostMetrics{1.5 + 2.5 + 0.1, nonnegative_int{1 + 2}},
ParallelLayerGuidObliviousMachineMapping{{
{
BinaryTreePath{{
Expand Down Expand Up @@ -266,7 +270,7 @@ TEST_SUITE(FF_TEST_SUITE) {
cache, context, problem_tree, full_machine_spec, constraints);
MachineMappingWithMemoryResult correct =
MachineMappingWithMemoryResult{{MachineMappingForSingleLayer{
OpCostMetrics{2.5, 2},
OpCostMetrics{2.5, nonnegative_int{2}},
ParallelLayerGuidObliviousMachineMapping{{
{
BinaryTreePath{{
Expand Down
20 changes: 10 additions & 10 deletions ...st/src/compiler/machine_mapping/memory_optimization/machine_mapping_result_with_memory.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,15 +53,15 @@ TEST_SUITE(FF_TEST_SUITE) {

OpCostMetrics cost1 = OpCostMetrics{
/*runtime=*/2.0,
/*memory=*/2,
/*memory=*/nonnegative_int{2},
};
OpCostMetrics cost2 = OpCostMetrics{
/*runtime=*/4.0,
/*memory=*/1,
/*memory=*/nonnegative_int{1},
};
OpCostMetrics cost3 = OpCostMetrics{
/*runtime=*/2.0,
/*memory=*/3,
/*memory=*/nonnegative_int{3},
};

MachineMappingForSingleLayer mm1 = MachineMappingForSingleLayer{
Expand Down Expand Up @@ -183,7 +183,7 @@ TEST_SUITE(FF_TEST_SUITE) {

OpCostMetrics pre_cost = OpCostMetrics{
/*runtime=*/2.0,
/*memory=*/2,
/*memory=*/nonnegative_int{2},
};
MachineMappingWithMemoryResult pre = MachineMappingWithMemoryResult{{
MachineMappingForSingleLayer{
Expand All @@ -209,7 +209,7 @@ TEST_SUITE(FF_TEST_SUITE) {

OpCostMetrics post_cost = OpCostMetrics{
/*runtime=*/4.0,
/*memory=*/1,
/*memory=*/nonnegative_int{1},
};

MachineMappingWithMemoryResult post = MachineMappingWithMemoryResult{{
Expand Down Expand Up @@ -378,7 +378,7 @@ TEST_SUITE(FF_TEST_SUITE) {

OpCostMetrics lhs_cost = OpCostMetrics{
/*runtime=*/2.0,
/*memory=*/2,
/*memory=*/nonnegative_int{2},
};
MachineMappingWithMemoryResult lhs = MachineMappingWithMemoryResult{{
MachineMappingForSingleLayer{
Expand All @@ -404,7 +404,7 @@ TEST_SUITE(FF_TEST_SUITE) {

OpCostMetrics rhs_cost = OpCostMetrics{
/*runtime=*/4.0,
/*memory=*/1,
/*memory=*/nonnegative_int{1},
};
MachineMappingWithMemoryResult rhs = MachineMappingWithMemoryResult{{
MachineMappingForSingleLayer{
Expand Down Expand Up @@ -519,15 +519,15 @@ TEST_SUITE(FF_TEST_SUITE) {

OpCostMetrics cost1 = OpCostMetrics{
/*runtime=*/2.0,
/*memory=*/2,
/*memory=*/nonnegative_int{2},
};
OpCostMetrics cost2 = OpCostMetrics{
/*runtime=*/4.0,
/*memory=*/1,
/*memory=*/nonnegative_int{1},
};
OpCostMetrics cost3 = OpCostMetrics{
/*runtime=*/2.0,
/*memory=*/3,
/*memory=*/nonnegative_int{3},
};

MachineMappingForSingleLayer mm1 = MachineMappingForSingleLayer{
Expand Down
2 changes: 2 additions & 0 deletions lib/utils/include/utils/nonnegative_int/nonnegative_int.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ class nonnegative_int {
friend bool operator!=(int const &lhs, nonnegative_int const &rhs);
friend bool operator>=(int const &lhs, nonnegative_int const &rhs);

nonnegative_int operator+(nonnegative_int const &other) const;

friend std::ostream &operator<<(std::ostream &os, nonnegative_int const &n);

friend int format_as(nonnegative_int const &);
Expand Down
4 changes: 4 additions & 0 deletions lib/utils/src/utils/nonnegative_int/nonnegative_int.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,10 @@ bool operator>=(int const &lhs, nonnegative_int const &rhs) {
return lhs >= rhs.value_;
}

nonnegative_int nonnegative_int::operator+(nonnegative_int const &other) const {
return nonnegative_int{this->value_ + other.value_};
}

std::ostream &operator<<(std::ostream &os, nonnegative_int const &n) {
os << n.value_;
return os;
Expand Down
9 changes: 9 additions & 0 deletions lib/utils/test/src/utils/nonnegative_int/nonnegative_int.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,15 @@ TEST_SUITE(FF_TEST_SUITE) {
}
}

TEST_CASE("nonnegative_int + operation") {
nonnegative_int nn_int_1a = nonnegative_int{1};
nonnegative_int nn_int_1b = nonnegative_int{1};
nonnegative_int nn_int_2 = nonnegative_int{2};
SUBCASE("LHS: nonnegative_int, RHS: nonnegative_int") {
CHECK(nn_int_1a + nn_int_1b == nn_int_2);
}
}

TEST_CASE("adl_serializer<nonnegative_int>") {
SUBCASE("to_json") {
nonnegative_int input = nonnegative_int{5};
Expand Down