Supplement ND SBP signatures for reshape op

oneflow/core/framework/sbp_infer_util.cpp

@@ -787,85 +787,67 @@ void DfsGetNdSbpSignature(NdSbpSignature& nd_sbp_sig, int32_t depth, int32_t dim
   }
 }
 
-// return -1 indicate lhs < rhs, 1 indicate lhs > rhs, 0 indicate lhs == rhs
-int CompareNdSbp(const NdSbp& lhs_nd_sbp, const NdSbp& rhs_nd_sbp) {
-  CHECK_EQ(lhs_nd_sbp.sbp_parallel_size(), rhs_nd_sbp.sbp_parallel_size());
-  // max split axis = 8, + B + P
-  const size_t kMaxSplitAxis = 8;
-  const size_t kCarryDigit = kMaxSplitAxis + 2;
-  auto Mesure = [](const NdSbp& nd_sbp) -> size_t {
-    size_t value = 0;
-    for (int i = 0; i < nd_sbp.sbp_parallel_size(); ++i) {
-      size_t cur_dim_value = 0;
-      const auto& sbp = nd_sbp.sbp_parallel(i);
-      if (sbp.has_split_parallel()) {
-        CHECK_LT(sbp.split_parallel().axis(), kMaxSplitAxis);
-        cur_dim_value = sbp.split_parallel().axis();
-      } else if (sbp.has_broadcast_parallel()) {
-        cur_dim_value = kMaxSplitAxis;
-      } else if (sbp.has_partial_sum_parallel()) {
-        cur_dim_value = kMaxSplitAxis + 1;
-      } else {
-        UNIMPLEMENTED();
-      }
-      value += cur_dim_value * std::pow(kCarryDigit, (nd_sbp.sbp_parallel_size() - i - 1));
-    }
-    return value;
-  };
-  size_t lhs_value = Mesure(lhs_nd_sbp);
-  size_t rhs_value = Mesure(rhs_nd_sbp);
+namespace {
 
-  if (lhs_value < rhs_value) {
-    return -1;
-  } else if (lhs_value > rhs_value) {
-    return 1;
-  } else {
-    return 0;
+// give a mesure value for NdSbp for sorting
+size_t MesureNdSbp(const NdSbp& nd_sbp) {
+  // max split axis = 8 (start from 1), + B + P = 10
+  constexpr size_t kMaxSplitAxis = 8;
+  constexpr size_t kCarryDigit = kMaxSplitAxis + 3;
+  size_t value = 0;
+  for (int i = 0; i < nd_sbp.sbp_parallel_size(); ++i) {
+    size_t cur_dim_value = 0;
+    const auto& sbp = nd_sbp.sbp_parallel(i);
+    if (sbp.has_split_parallel()) {
+      CHECK_LT(sbp.split_parallel().axis(), kMaxSplitAxis);
+      // from 1 to 8
+      cur_dim_value = sbp.split_parallel().axis() + 1;
+    } else if (sbp.has_broadcast_parallel()) {
+      // 9
+      cur_dim_value = kMaxSplitAxis + 1;
+    } else if (sbp.has_partial_sum_parallel()) {
+      // 10
+      cur_dim_value = kMaxSplitAxis + 2;
+    } else {
+      UNIMPLEMENTED();
+    }
+    value = value * kCarryDigit + cur_dim_value;
   }
+  return value;
 }
 
-// return -1 indicate lhs < rhs, 1 indicate lhs > rhs, 0 indicate lhs == rhs
-int CompareNdSbpSignature(const NdSbpSignature& lhs_nd_sbp_sig,
-                          const NdSbpSignature& rhs_nd_sbp_sig) {
-  CHECK_EQ(lhs_nd_sbp_sig.bn_in_op2nd_sbp_size(), rhs_nd_sbp_sig.bn_in_op2nd_sbp_size());
-  for (const auto& lhs_bn_nd_sbp : lhs_nd_sbp_sig.bn_in_op2nd_sbp()) {
-    const auto& bn = lhs_bn_nd_sbp.first;
-    auto rhs_bn_nd_sbp_it = rhs_nd_sbp_sig.bn_in_op2nd_sbp().find(bn);
-    CHECK(rhs_bn_nd_sbp_it != rhs_nd_sbp_sig.bn_in_op2nd_sbp().end());
-    const auto& lhs_nd_sbp = lhs_bn_nd_sbp.second;
-    const auto& rhs_nd_sbp = rhs_bn_nd_sbp_it->second;
-    int cmp_ret = CompareNdSbp(lhs_nd_sbp, rhs_nd_sbp);
-    if (cmp_ret != 0) { return cmp_ret; }
+size_t MesureNdSbpSignature(const NdSbpSignature& nd_sbp_sig, const std::vector<std::string>& bns) {
+  // big enough for 2d-sbp signatrue set
+  // if want to extend to 3d-sbp, consider increase to 170
+  constexpr size_t kCarryDigit = 97;
+  size_t value = 0;
+  for (size_t i = 0; i < bns.size(); ++i) {
+    auto nd_sbp_it = nd_sbp_sig.bn_in_op2nd_sbp().find(bns[i]);
+    CHECK(nd_sbp_it != nd_sbp_sig.bn_in_op2nd_sbp().end())
+        << "can't find bn (" << bns[i] << ") in " << PbMessage2TxtString(nd_sbp_sig);
+    size_t cur_arg_value = MesureNdSbp(nd_sbp_it->second);
+    CHECK_LE(value + cur_arg_value / kCarryDigit, std::numeric_limits<size_t>::max() / kCarryDigit);
+    value = value * kCarryDigit + cur_arg_value;
   }
-  return 0;
+  return value;
 }
 
-void DeduplicateNdSbpSignatureList(std::vector<NdSbpSignature>* nd_sbp_sig_list) {
-  auto CompareIndices = [nd_sbp_sig_list](size_t lhs, size_t rhs) -> bool {
-    int cmp_ret = CompareNdSbpSignature(nd_sbp_sig_list->at(lhs), nd_sbp_sig_list->at(rhs));
-    if (cmp_ret == 0) {
-      return true;
-    } else if (cmp_ret == -1) {
-      return true;
-    } else if (cmp_ret == 1) {
-      return false;
-    } else {
-      UNIMPLEMENTED();
+}  // namespace
+
+void DeduplicateNdSbpSignatureList(std::vector<NdSbpSignature>* nd_sbp_sig_list,
+                                   const std::vector<std::string>& bns) {
+  if (bns.size() > 8) { return; }
+  std::map<size_t, NdSbpSignature> value2nd_sbp_sig;
+  for (auto& nd_sbp_sig : *nd_sbp_sig_list) {
+    size_t order_value = MesureNdSbpSignature(nd_sbp_sig, bns);
+    if (value2nd_sbp_sig.find(order_value) == value2nd_sbp_sig.end()) {
+      value2nd_sbp_sig.emplace(order_value, std::move(nd_sbp_sig));
     }
-  };
-  std::vector<size_t> indices(nd_sbp_sig_list->size());
-  std::iota(indices.begin(), indices.end(), 0);
-  std::sort(indices.begin(), indices.end(), CompareIndices);
-  auto new_end =
-      std::unique(indices.begin(), indices.end(), [nd_sbp_sig_list](size_t lhs, size_t rhs) {
-        return nd_sbp_sig_list->at(lhs) == nd_sbp_sig_list->at(rhs);
-      });
-  indices.erase(new_end, indices.end());
-  std::vector<NdSbpSignature> new_nd_sbp_sig_list;
-  for (size_t index : indices) {
-    new_nd_sbp_sig_list.emplace_back(std::move(nd_sbp_sig_list->at(index)));
   }
-  *nd_sbp_sig_list = std::move(new_nd_sbp_sig_list);
+  nd_sbp_sig_list->clear();
+  for (auto& nd_sbp_pair : value2nd_sbp_sig) {
+    nd_sbp_sig_list->emplace_back(std::move(nd_sbp_pair.second));
+  }
 }
 
 // Compute storage per device for given NdSbp

oneflow/core/framework/sbp_infer_util.h

@@ -77,11 +77,8 @@ void DfsGetNdSbpSignature(NdSbpSignature& nd_sbp_sig, int32_t depth, int32_t dim
                           const HashMap<int32_t, SbpSignatureList>& hierarchy_value2sbp_sig_list,
                           std::vector<NdSbpSignature>* nd_sbp_sig_list);
 
-// return -1 indicate lhs < rhs, 1 indicate lhs > rhs, 0 indicate lhs == rhs
-int CompareNdSbp(const NdSbp& lhs_nd_sbp, const NdSbp& rhs_nd_sbp);
-int CompareNdSbpSignature(const NdSbpSignature& lhs_nd_sbp_sig,
-                          const NdSbpSignature& rhs_nd_sbp_sig);
-void DeduplicateNdSbpSignatureList(std::vector<NdSbpSignature>* nd_sbp_sig_list);
+void DeduplicateNdSbpSignatureList(std::vector<NdSbpSignature>* nd_sbp_sig_list,
+                                   const std::vector<std::string>& bns);
 
 // Compute storage for given NdSbp
 double Storage4NdSbp(const NdSbp& nd_sbp, Shape& logical_shape, const Shape& parallel_hierarchy);

oneflow/core/framework/sbp_infer_util_test.cpp

@@ -71,6 +71,7 @@ bool ParseNdSbpSignatureFromString(const std::string& nd_sbp_signature_str,
       if (pos < nd_sbp_signature_str.size() && nd_sbp_signature_str[pos] == '>') {
         // in args parsing has finished, parse out args
         arg_name = "out";
+        nd_sbp_group_id = 0;
         // skip '>' in substr '->'
         pos++;
         continue;
@@ -85,51 +86,56 @@ bool ParseNdSbpSignatureFromString(const std::string& nd_sbp_signature_str,
   return true;
 }
 
-void TestCompareNdSbpSignature(const std::vector<std::string>& nd_sbp_signature_str_list) {
+void TestDeduplicateNdSbpSignature(const std::vector<std::string>& nd_sbp_signature_str_list,
+                                   const std::vector<std::string>& bns) {
   std::vector<NdSbpSignature> nd_sbp_sig_list;
+  nd_sbp_sig_list.reserve(nd_sbp_signature_str_list.size());
   for (const auto& nd_sbp_signature_str : nd_sbp_signature_str_list) {
     nd_sbp_sig_list.emplace_back();
     ASSERT_TRUE(ParseNdSbpSignatureFromString(nd_sbp_signature_str, nd_sbp_sig_list.back()));
   }
-  for (const auto& lhs_nd_sbp_sig : nd_sbp_sig_list) {
-    for (const auto& rhs_nd_sbp_sig : nd_sbp_sig_list) {
-      int cmp_ret1 = CompareNdSbpSignature(lhs_nd_sbp_sig, rhs_nd_sbp_sig);
-      int cmp_ret2 = CompareNdSbpSignature(rhs_nd_sbp_sig, lhs_nd_sbp_sig);
-      ASSERT_TRUE(cmp_ret1 + cmp_ret2 == 0);
-    }
-  }
+  std::random_device rd;
+  std::mt19937 gen(rd());
+  std::shuffle(nd_sbp_sig_list.begin(), nd_sbp_sig_list.end(), gen);
+  nd_sbp_sig_list.reserve(nd_sbp_sig_list.size() + nd_sbp_sig_list.size() / 2);
+  std::copy_n(nd_sbp_sig_list.begin(), nd_sbp_sig_list.size() / 2,
+              std::back_inserter(nd_sbp_sig_list));
+  std::shuffle(nd_sbp_sig_list.begin(), nd_sbp_sig_list.end(), gen);
+  DeduplicateNdSbpSignatureList(&nd_sbp_sig_list, bns);
 }
 
 }  // namespace
 
-TEST(SbpInferUtil, CompareNdSbpSignature) {
-  TestCompareNdSbpSignature({
-      "(S(0), S(0)) -> (S(0), S(0))",
-      "(S(0), S(1)) -> (S(0), S(1))",
-      "(S(0), S(3)) -> (S(0), S(2))",
-      "(S(0), B) -> (S(0), B)",
-      "(S(0), P) -> (S(0), P)",
-      "(S(1), S(0)) -> (S(1), S(0))",
-      "(S(1), S(1)) -> (S(1), S(1))",
-      "(S(1), S(3)) -> (S(1), S(2))",
-      "(S(1), B) -> (S(1), B)",
-      "(S(1), P) -> (S(1), P)",
-      "(S(3), S(0)) -> (S(2), S(0))",
-      "(S(3), S(1)) -> (S(2), S(1))",
-      "(S(3), S(3)) -> (S(2), S(2))",
-      "(S(3), B) -> (S(2), B)",
-      "(S(3), P) -> (S(2), P)",
-      "(B, S(0)) -> (B, S(0))",
-      "(B, S(1)) -> (B, S(1))",
-      "(B, S(3)) -> (B, S(2))",
-      "(B, B) -> (B, B)",
-      "(B, P) -> (B, P)",
-      "(P, S(0)) -> (P, S(0))",
-      "(P, S(1)) -> (P, S(1))",
-      "(P, S(3)) -> (P, S(2))",
-      "(P, B) -> (P, B)",
-      "(P, P) -> (P, P)",
-  });
+TEST(SbpInferUtil, DeduplicateNdSbpSignatureList) {
+  TestDeduplicateNdSbpSignature(
+      {
+          "(S(0), S(0)) -> (S(0), S(0))",
+          "(S(0), S(1)) -> (S(0), S(1))",
+          "(S(0), S(3)) -> (S(0), S(2))",
+          "(S(0), B) -> (S(0), B)",
+          "(S(0), P) -> (S(0), P)",
+          "(S(1), S(0)) -> (S(1), S(0))",
+          "(S(1), S(1)) -> (S(1), S(1))",
+          "(S(1), S(3)) -> (S(1), S(2))",
+          "(S(1), B) -> (S(1), B)",
+          "(S(1), P) -> (S(1), P)",
+          "(S(3), S(0)) -> (S(2), S(0))",
+          "(S(3), S(1)) -> (S(2), S(1))",
+          "(S(3), S(3)) -> (S(2), S(2))",
+          "(S(3), B) -> (S(2), B)",
+          "(S(3), P) -> (S(2), P)",
+          "(B, S(0)) -> (B, S(0))",
+          "(B, S(1)) -> (B, S(1))",
+          "(B, S(3)) -> (B, S(2))",
+          "(B, B) -> (B, B)",
+          "(B, P) -> (B, P)",
+          "(P, S(0)) -> (P, S(0))",
+          "(P, S(1)) -> (P, S(1))",
+          "(P, S(3)) -> (P, S(2))",
+          "(P, B) -> (P, B)",
+          "(P, P) -> (P, P)",
+      },
+      {"in_0", "out_0"});
 }
 
 }  // namespace test

oneflow/user/ops/reshape_op.cpp

@@ -61,7 +61,7 @@ namespace oneflow {
     }
   }
 
-  DeduplicateNdSbpSignatureList(nd_sbp_sig_list);
+  DeduplicateNdSbpSignatureList(nd_sbp_sig_list, {"in_0", "out_0"});
   return Maybe<void>::Ok();
 }