Function CanMergeSplit

oneflow/core/framework/sbp_infer_util.cpp:397–404

We can not just simply merging two same split For example, shape = [6], we are trying to merge [2, 2]: (S0, S0) -> [4]: S0 For each rank, [4]: S0 has number of data: 2, 2, 1, 1 For each rank, [2]: S0 has number of data: 3, 3 For each rank, [2, 2]: (S0, S0) has number of data: 2, 1, 2, 1 Thus {[2, 2]: (S0, S0)} != {[4]: S0} for shape [6] However {[2, 2]: (S0, S0)} == {[4]: S0} for shape [4], [5], [

// More specifically, {[a, b]: (Si, Si)} == {[a*b]: Si} if and only if
// shape value % (a * b) == 0, 1, a*b - 1
bool CanMergeSplit(int32_t shape_value, int32_t merged_split_hierarchy_value) {
  int32_t remainder = shape_value % merged_split_hierarchy_value;
  if (remainder <= 1 || remainder == merged_split_hierarchy_value - 1) {
    return true;
  } else {
    return false;
  }
}

}  // namespace