Function InOutParallelDimReduce

oneflow/core/framework/sbp_infer_util.cpp:557–605

}

void InOutParallelDimReduce(const Shape& in_hierarchy, const Shape& out_hierarchy,
                            const NdSbp& in_nd_sbp, const NdSbp& out_nd_sbp,
                            Shape* reduced_in_hierarchy, Shape* reduced_out_hierarchy,
                            NdSbp* reduced_in_nd_sbp, NdSbp* reduced_out_nd_sbp,
                            const Shape& logical_shape) {
  if (in_hierarchy == out_hierarchy) {
    // [2, 4]: (S0, S0) -> [2, 4]: (S0, S1)
    NdSbpsDimReduce(in_hierarchy, {&in_nd_sbp, &out_nd_sbp}, reduced_in_hierarchy,
                    {reduced_in_nd_sbp, reduced_out_nd_sbp}, logical_shape);
    *reduced_out_hierarchy = *reduced_in_hierarchy;
  } else {
    // [2, 4]: (S0, S0) -> [4, 2]: (S0, S1)
    // [2, 4]: (S0, S0) -> [3, 3]: (S0, S1)
    NdSbpDimReduce(in_hierarchy, in_nd_sbp, reduced_in_hierarchy, reduced_in_nd_sbp, logical_shape);
    NdSbpDimReduce(out_hierarchy, out_nd_sbp, reduced_out_hierarchy, reduced_out_nd_sbp,
                   logical_shape);

    // Sbp of 3d or higher dimension would use general basic communication
    // Only looks at 1d to 2d or 2d to 1d
    if (reduced_in_hierarchy->NumAxes() + reduced_out_hierarchy->NumAxes() == 3
        && reduced_in_hierarchy->elem_cnt() == reduced_out_hierarchy->elem_cnt()) {
      if (reduced_in_hierarchy->NumAxes() == 1) {
        // [8]: S0 -> [4, 2]: (S0, S1)
        // [8]: B -> [2, 4]: (S0, S1)
        const auto& in_sbp_parallel = reduced_in_nd_sbp->sbp_parallel(0);
        if (!in_sbp_parallel.has_split_parallel()
            || CanMergeSplit(logical_shape.At(in_sbp_parallel.split_parallel().axis()),
                             reduced_in_hierarchy->elem_cnt())) {
          // Change [8]: S0 -> [4, 2]: (S0, S1) to [4, 2]: (S0, S0) -> [4, 2]: (S0, S1)
          // Change [8]: B -> [2, 4]: (S0, S1) to [2, 4]: (B, B) -> [2, 4]: (S0, S1)
          *reduced_in_nd_sbp->add_sbp_parallel() = in_sbp_parallel;
          *reduced_in_hierarchy = *reduced_out_hierarchy;
        }
      } else {
        // [2, 3]: (S0, P) -> [6]: S0
        // [3, 4]: (B, S1) -> [12]: B
        const auto& out_sbp_parallel = reduced_out_nd_sbp->sbp_parallel(0);
        if (!out_sbp_parallel.has_split_parallel()
            || CanMergeSplit(logical_shape.At(out_sbp_parallel.split_parallel().axis()),
                             reduced_out_hierarchy->elem_cnt())) {
          // Change [2, 3]: (S0, P) -> [6]: S0 to [2, 3]: (S0, P) -> [2, 3]: (S0, S0)
          // Change [3, 4]: (B, S1) -> [12]: B to [3, 4]: (B, S1) -> [3, 4]: (B, B)
          *reduced_out_nd_sbp->add_sbp_parallel() = out_sbp_parallel;
          *reduced_out_hierarchy = *reduced_in_hierarchy;
        }
      }
    }
  }
}

void InOutParallelDimReduce(const ParallelDesc& in_parallel_desc,
                            const ParallelDesc& out_parallel_desc, const NdSbp& in_nd_sbp,