Oneflow 0.8 for DCU

oneflow/core/boxing/eager_boxing_interpreter.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include <typeinfo>
+#include "oneflow/core/common/container_util.h"
+#include "oneflow/core/common/registry_error.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/framework/tensor_rpc_util.h"
+#include "oneflow/core/framework/to_string.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter_mgr.h"
+#include "oneflow/core/framework/nd_sbp.h"
+
+namespace oneflow {
+
+namespace {
+Maybe<void> CheckEagerBoxingDataType(DataType val) {
+  CHECK_OR_RETURN(val != DataType::kTensorBuffer && val != DataType::kOFRecord)
+      << Error::RuntimeError() << "invalid boxing data type " << ToString(val);
+  return Maybe<void>::Ok();
+}
+}  // namespace
+
+Maybe<one::Tensor> EagerBoxingInterpreter::Interpret(const std::shared_ptr<one::Tensor>& input,
+                                                     Symbol<NdSbp> in_nd_sbp,
+                                                     Symbol<NdSbp> out_nd_sbp,
+                                                     Symbol<ParallelDesc> in_parallel_desc,
+                                                     Symbol<ParallelDesc> out_parallel_desc) const {
+  JUST(CheckEagerBoxingDataType(input->dtype()->data_type()));
+  DisableCheckConsistentTensorMetaScope disable_meta_check;
+  const auto& tensor =
+      JUST(InterpretImpl(input, in_nd_sbp, out_nd_sbp, in_parallel_desc, out_parallel_desc));
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_nd_sbp == out_nd_sbp)
+      << Error::RuntimeError() << "The sbp of output tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the output sbp (" << NdSbpToString(out_nd_sbp) << ")";
+  CHECK_OR_RETURN(tensor_placement == out_parallel_desc)
+      << Error::RuntimeError() << "The placement of output tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the output placement ("
+      << *JUST(PlacementToString(out_parallel_desc)) << ")";
+  return tensor;
+}
+
+namespace {
+
+HashMap<std::string, BoxingCheckerT>* MutName2BoxingChecker() {
+  static HashMap<std::string, BoxingCheckerT> map;
+  return &map;
+}
+
+HashMap<std::string, BoxingFunctionT>* MutName2BoxingFunction() {
+  static HashMap<std::string, BoxingFunctionT> map;
+  return &map;
+}
+
+Maybe<BoxingFunctionT> RawGetBoxingFunction(const std::string& method_name, Symbol<PlacedNdSbp> in,
+                                            Symbol<PlacedNdSbp> out, const Shape& logical_shape) {
+  const auto& Checker =
+      JUST_MSG(MapAt(*MutName2BoxingChecker(), method_name),
+               std::stringstream() << "boxing checker not found. checker_name: " << method_name);
+  JUST(Checker(in, out, logical_shape));
+  return JUST_MSG(MapAt(*MutName2BoxingFunction(), method_name),
+                  std::stringstream()
+                      << "boxing function not found. function_name: " << method_name);
+}
+
+}  // namespace
+
+Maybe<BoxingFunctionT> GetBoxingFunction(const std::string& method_name, Symbol<PlacedNdSbp> in,
+                                         Symbol<PlacedNdSbp> out, const Shape& logical_shape) {
+  return DECORATE(&RawGetBoxingFunction, ThreadLocalCachedCopiable)(method_name, in, out,
+                                                                    logical_shape);
+}
+
+void RegisterBoxingFunction(const std::string& method_name, const BoxingCheckerT& Checker,
+                            const BoxingFunctionT& BoxingFunction) {
+  CatchRegistryError([&]() -> Maybe<void> {
+    CHECK_OR_RETURN(MutName2BoxingChecker()->emplace(method_name, Checker).second)
+        << Error::RuntimeError() << "register boxing checker failed: " << method_name;
+    CHECK_OR_RETURN(MutName2BoxingFunction()->emplace(method_name, BoxingFunction).second)
+        << Error::RuntimeError() << "register boxing function failed: " << method_name;
+    return Maybe<void>::Ok();
+  });
+}
+
+Maybe<BoxingInterpreterStatus> AtomicBoxingExpr::Check(Symbol<PlacedNdSbp> in,
+                                                       Symbol<PlacedNdSbp> out,
+                                                       const Shape& logical_shape) const {
+  const auto& Checker =
+      JUST_MSG(MapAt(*MutName2BoxingChecker(), boxing_name_),
+               std::stringstream() << "boxing checker not found. checker_name: " << boxing_name_);
+  JUST(Checker(in, out, logical_shape));
+  return MakeBoxingInterpreterStatus(boxing_name_, logical_shape, in, out);
+}
+
+Maybe<BoxingFunctionT> AtomicBoxingExpr::GetBoxingFunction(Symbol<PlacedNdSbp> in,
+                                                           Symbol<PlacedNdSbp> out,
+                                                           const Shape& logical_shape) const {
+  return DECORATE(&RawGetBoxingFunction, ThreadLocalCachedCopiable)(boxing_name_, in, out,
+                                                                    logical_shape);
+}
+
+Maybe<BoxingInterpreterStatus> DivideAndConquerBoxingExpr::Check(Symbol<PlacedNdSbp> in,
+                                                                 Symbol<PlacedNdSbp> out,
+                                                                 const Shape& logical_shape) const {
+  const auto& middle = JUST((*boxing_dividor_)(in, out));
+  const auto& lhs_status = JUST(lhs_conquer_->Check(in, middle, logical_shape));
+  const auto& rhs_status = JUST(rhs_conquer_->Check(middle, out, logical_shape));
+  return MakeComposedBoxingInterpreterStatus(lhs_status, rhs_status);
+}
+
+Maybe<BoxingFunctionT> DivideAndConquerBoxingExpr::GetBoxingFunction(
+    Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out, const Shape& logical_shape) const {
+  const auto& middle = JUST((*boxing_dividor_)(in, out));
+  const auto& lhs_boxing_func = JUST(lhs_conquer_->GetBoxingFunction(in, middle, logical_shape));
+  const auto& rhs_boxing_func = JUST(rhs_conquer_->GetBoxingFunction(middle, out, logical_shape));
+  BoxingFunctionT boxing_function =
+      [lhs_boxing_func, rhs_boxing_func, middle, in, out, &logical_shape](
+          const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> arg_in,
+          Symbol<PlacedNdSbp> arg_out) -> Maybe<one::Tensor> {
+    // Always true, if check failed, there is a bug in oneflow needed to be resolved.
+    CHECK_OR_RETURN(in == arg_in) << Error::RuntimeError() << "The placement ("
+                                  << *JUST(PlacementToString(arg_in->placement())) << ") and sbp ("
+                                  << NdSbpToString(in->nd_sbp())
+                                  << ") of input tensor must match the placement ("
+                                  << *JUST(PlacementToString(in->placement())) << ") and sbp ("
+                                  << NdSbpToString(arg_in->nd_sbp())
+                                  << ") used for get this boxing function! Please submit an issue "
+                                     "in `https://github.com/Oneflow-Inc/oneflow/issues` "
+                                     "and we will fix it as soon as possible";
+    CHECK_OR_RETURN(logical_shape == *tensor->shape())
+        << Error::RuntimeError() << "The logical_shape " << tensor->shape()->ToString()
+        << " of input tensor must match the logical_shape " << logical_shape.ToString()
+        << " used for get this boxing function! Please submit an issue in "
+           "`https://github.com/Oneflow-Inc/oneflow/issues` and we will fix it "
+           "as soon as possible";
+    CHECK_OR_RETURN(out == arg_out)
+        << Error::RuntimeError() << "The placement ("
+        << *JUST(PlacementToString(arg_out->placement())) << ") and sbp ("
+        << NdSbpToString(arg_out->nd_sbp()) << ") of output tensor must match the placement ("
+        << *JUST(PlacementToString(out->placement())) << ") and sbp ("
+        << NdSbpToString(out->nd_sbp())
+        << ") used for get this boxing function! Please submit "
+           "an issue in `https://github.com/Oneflow-Inc/oneflow/issues` and we will fix it "
+           "as soon as possible";
+    const auto& middle_tensor = JUST((*lhs_boxing_func)(tensor, in, middle));
+    return JUST((*rhs_boxing_func)(middle_tensor, middle, out));
+  };
+  return boxing_function;
+}
+
+Maybe<BoxingInterpreterStatus> OrBoxingExpr::Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                                   const Shape& logical_shape) const {
+  const auto& lhs_status = TRY(lhs_boxing_->Check(in, out, logical_shape));
+  if (lhs_status.IsOk()) { return lhs_status; }
+  return rhs_boxing_->Check(in, out, logical_shape);
+}
+
+Maybe<BoxingFunctionT> OrBoxingExpr::GetBoxingFunction(Symbol<PlacedNdSbp> in,
+                                                       Symbol<PlacedNdSbp> out,
+                                                       const Shape& logical_shape) const {
+  if (lhs_boxing_->Check(in, out, logical_shape).IsOk()) {
+    return lhs_boxing_->GetBoxingFunction(in, out, logical_shape);
+  }
+  JUST(rhs_boxing_->Check(in, out, logical_shape));
+  return rhs_boxing_->GetBoxingFunction(in, out, logical_shape);
+}
+
+Maybe<BoxingExprIf> BoxingExpr(const std::string& boxing_name) {
+  JUST(MapAt(*MutName2BoxingChecker(), boxing_name));
+  auto boxing_expr = std::make_unique<AtomicBoxingExpr>(boxing_name);
+  return std::shared_ptr<BoxingExprIf>(std::move(boxing_expr));
+}
+
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::string& lhs_conquer, const std::string& rhs_conquer) {
+  return BoxingExpr(boxing_dividor, JUST(BoxingExpr(lhs_conquer)), JUST(BoxingExpr(rhs_conquer)));
+}
+
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::shared_ptr<BoxingExprIf>& lhs_conquer,
+                               const std::string& rhs_conquer) {
+  return BoxingExpr(boxing_dividor, lhs_conquer, JUST(BoxingExpr(rhs_conquer)));
+}
+
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::string& lhs_conquer,
+                               const std::shared_ptr<BoxingExprIf>& rhs_conquer) {
+  return BoxingExpr(boxing_dividor, JUST(BoxingExpr(lhs_conquer)), rhs_conquer);
+}
+
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::shared_ptr<BoxingExprIf>& lhs_conquer,
+                               const std::shared_ptr<BoxingExprIf>& rhs_conquer) {
+  auto divide_and_conquer =
+      std::make_unique<DivideAndConquerBoxingExpr>(boxing_dividor, lhs_conquer, rhs_conquer);
+  return std::shared_ptr<BoxingExprIf>(std::move(divide_and_conquer));
+}
+
+std::shared_ptr<BoxingExprIf> operator|(const std::shared_ptr<BoxingExprIf>& lhs_boxing,
+                                        const std::shared_ptr<BoxingExprIf>& rhs_boxing) {
+  auto or_boxing = std::make_unique<OrBoxingExpr>(lhs_boxing, rhs_boxing);
+  return std::shared_ptr<BoxingExprIf>(std::move(or_boxing));
+}
+
+Maybe<BoxingExprIf> OptionalBoxing(const std::string& boxing_mame) {
+  return JUST(BoxingExpr(boxing_mame)) | JUST(BoxingExpr("identity"));
+}
+
+}  // namespace oneflow

oneflow/core/boxing/eager_boxing_interpreter.h

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_BOXING_EAGER_BOXING_INTERPRETER_H_
+#define ONEFLOW_CORE_BOXING_EAGER_BOXING_INTERPRETER_H_
+
+#include "oneflow/core/common/symbol.h"
+#include "oneflow/core/framework/tensor_tuple.h"
+#include "oneflow/core/boxing/boxing_dividor.h"
+#include "oneflow/core/framework/tensor.h"
+#include "oneflow/core/framework/placed_nd_sbp.h"
+#include "oneflow/core/job/parallel_desc.h"
+#include "oneflow/core/job/sbp_parallel.h"
+#include "oneflow/core/boxing/boxing_interpreter_status.h"
+
+namespace oneflow {
+
+class EagerBoxingInterpreter {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(EagerBoxingInterpreter);
+  EagerBoxingInterpreter() = default;
+  virtual ~EagerBoxingInterpreter() = default;
+
+  Maybe<one::Tensor> Interpret(const std::shared_ptr<one::Tensor>& input, Symbol<NdSbp> in_nd_sbp,
+                               Symbol<NdSbp> out_nd_sbp, Symbol<ParallelDesc> in_parallel_desc,
+                               Symbol<ParallelDesc> out_parallel_desc) const;
+  virtual Maybe<BoxingInterpreterStatus> boxing_interpreter_status() const = 0;
+
+ protected:
+  virtual Maybe<one::Tensor> InterpretImpl(const std::shared_ptr<one::Tensor>& input,
+                                           Symbol<NdSbp> in_nd_sbp, Symbol<NdSbp> out_nd_sbp,
+                                           Symbol<ParallelDesc> in_parallel_desc,
+                                           Symbol<ParallelDesc> out_parallel_desc) const = 0;
+};
+
+using BoxingCheckerT = std::function<Maybe<void>(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                                 const Shape& logical_shape)>;
+using BoxingFunctionT = std::function<Maybe<one::Tensor>(
+    const std::shared_ptr<one::Tensor>& input, Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out)>;
+
+Maybe<BoxingFunctionT> GetBoxingFunction(const std::string& method_name, Symbol<PlacedNdSbp> in,
+                                         Symbol<PlacedNdSbp> out, const Shape& logical_shape);
+
+void RegisterBoxingFunction(const std::string& method_name, const BoxingCheckerT& Check,
+                            const BoxingFunctionT& BoxingFunction);
+
+inline void RegisterBoxingFunction(
+    const std::string& method_name,
+    const std::pair<BoxingCheckerT, BoxingFunctionT>& CheckAndBoxing) {
+  RegisterBoxingFunction(method_name, CheckAndBoxing.first, CheckAndBoxing.second);
+}
+
+class NaiveEagerBoxingInterpreter : public EagerBoxingInterpreter {
+ public:
+  explicit NaiveEagerBoxingInterpreter(
+      const std::shared_ptr<BoxingFunctionT>& boxing_function,
+      const std::shared_ptr<BoxingInterpreterStatus>& boxing_interpreter_status)
+      : boxing_function_(boxing_function), boxing_interpreter_status_(boxing_interpreter_status) {}
+  NaiveEagerBoxingInterpreter(const NaiveEagerBoxingInterpreter&) = delete;
+  NaiveEagerBoxingInterpreter(NaiveEagerBoxingInterpreter&&) = delete;
+  ~NaiveEagerBoxingInterpreter() override = default;
+
+  Maybe<BoxingInterpreterStatus> boxing_interpreter_status() const override {
+    return boxing_interpreter_status_;
+  }
+
+ private:
+  Maybe<one::Tensor> InterpretImpl(const std::shared_ptr<one::Tensor>& input,
+                                   Symbol<NdSbp> in_nd_sbp, Symbol<NdSbp> out_nd_sbp,
+                                   Symbol<ParallelDesc> in_parallel_desc,
+                                   Symbol<ParallelDesc> out_parallel_desc) const override {
+    const auto& in_placed_nd_sbp = JUST(PlacedNdSbp::New(in_nd_sbp, in_parallel_desc));
+    const auto& out_placed_nd_sbp = JUST(PlacedNdSbp::New(out_nd_sbp, out_parallel_desc));
+    return JUST((*boxing_function_)(input, in_placed_nd_sbp, out_placed_nd_sbp));
+  }
+
+  const std::shared_ptr<BoxingFunctionT> boxing_function_;
+  const std::shared_ptr<BoxingInterpreterStatus> boxing_interpreter_status_;
+};
+
+class BoxingExprIf {
+ public:
+  BoxingExprIf(const BoxingExprIf&) = default;
+  BoxingExprIf(BoxingExprIf&&) = default;
+  virtual ~BoxingExprIf() = default;
+
+  virtual Maybe<BoxingInterpreterStatus> Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                               const Shape& logical_shape) const = 0;
+  virtual Maybe<BoxingFunctionT> GetBoxingFunction(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                                   const Shape& logical_shape) const = 0;
+
+ protected:
+  BoxingExprIf() = default;
+};
+
+class AtomicBoxingExpr final : public BoxingExprIf {
+ public:
+  AtomicBoxingExpr(const AtomicBoxingExpr&) = delete;
+  AtomicBoxingExpr(AtomicBoxingExpr&&) = delete;
+  ~AtomicBoxingExpr() override = default;
+
+  explicit AtomicBoxingExpr(const std::string& boxing_name)
+      : BoxingExprIf(), boxing_name_(boxing_name) {}
+
+  Maybe<BoxingInterpreterStatus> Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                       const Shape& logical_shape) const override;
+  Maybe<BoxingFunctionT> GetBoxingFunction(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                           const Shape& logical_shape) const override;
+
+ private:
+  const std::string boxing_name_;
+};
+
+class DivideAndConquerBoxingExpr final : public BoxingExprIf {
+ public:
+  DivideAndConquerBoxingExpr(const DivideAndConquerBoxingExpr&) = delete;
+  DivideAndConquerBoxingExpr(DivideAndConquerBoxingExpr&&) = delete;
+  ~DivideAndConquerBoxingExpr() override = default;
+
+  explicit DivideAndConquerBoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                                      const std::shared_ptr<BoxingExprIf>& lhs_conquer,
+                                      const std::shared_ptr<BoxingExprIf>& rhs_conquer)
+      : BoxingExprIf(),
+        boxing_dividor_(boxing_dividor),
+        lhs_conquer_(lhs_conquer),
+        rhs_conquer_(rhs_conquer) {}
+
+  Maybe<BoxingInterpreterStatus> Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                       const Shape& logical_shape) const override;
+  Maybe<BoxingFunctionT> GetBoxingFunction(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                           const Shape& logical_shape) const override;
+
+ private:
+  const std::shared_ptr<BoxingDividor> boxing_dividor_;
+  const std::shared_ptr<BoxingExprIf> lhs_conquer_;
+  const std::shared_ptr<BoxingExprIf> rhs_conquer_;
+};
+
+class OrBoxingExpr final : public BoxingExprIf {
+ public:
+  OrBoxingExpr(const OrBoxingExpr&) = delete;
+  OrBoxingExpr(OrBoxingExpr&&) = delete;
+  ~OrBoxingExpr() override = default;
+
+  explicit OrBoxingExpr(const std::shared_ptr<BoxingExprIf>& lhs_boxing,
+                        const std::shared_ptr<BoxingExprIf>& rhs_boxing)
+      : BoxingExprIf(), lhs_boxing_(lhs_boxing), rhs_boxing_(rhs_boxing) {}
+
+  Maybe<BoxingInterpreterStatus> Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                       const Shape& logical_shape) const override;
+  Maybe<BoxingFunctionT> GetBoxingFunction(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                           const Shape& logical_shape) const override;
+
+ private:
+  const std::shared_ptr<BoxingExprIf> lhs_boxing_;
+  const std::shared_ptr<BoxingExprIf> rhs_boxing_;
+};
+
+Maybe<BoxingExprIf> BoxingExpr(const std::string& boxing_name);
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::string& lhs_conquer, const std::string& rhs_conquer);
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::shared_ptr<BoxingExprIf>& lhs_conquer,
+                               const std::string& rhs_conquer);
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::string& lhs_conquer,
+                               const std::shared_ptr<BoxingExprIf>& rhs_conquer);
+Maybe<BoxingExprIf> BoxingExpr(const std::shared_ptr<BoxingDividor>& boxing_dividor,
+                               const std::shared_ptr<BoxingExprIf>& lhs_conquer,
+                               const std::shared_ptr<BoxingExprIf>& rhs_conquer);
+
+std::shared_ptr<BoxingExprIf> operator|(const std::shared_ptr<BoxingExprIf>& lhs_boxing,
+                                        const std::shared_ptr<BoxingExprIf>& rhs_boxing);
+
+Maybe<BoxingExprIf> OptionalBoxing(const std::string& boxing_mame);
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_BOXING_EAGER_BOXING_INTERPRETER_H_

oneflow/core/boxing/eager_boxing_interpreter_mgr.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include <utility>
+#include "oneflow/core/common/constant.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter_mgr.h"
+#include "oneflow/core/boxing/boxing_dividor_util.h"
+
+namespace oneflow {
+
+namespace {
+
+Maybe<bool> IgnoringDeviceTypeEqual(Symbol<ParallelDesc> lhs, Symbol<ParallelDesc> rhs) {
+  if (lhs == rhs) { return true; }
+  return lhs == JUST(ReplaceDeviceType(rhs, lhs->device_type()));
+}
+
+namespace {
+
+Maybe<BoxingExprIf> OptionalCudaCopy(const std::shared_ptr<BoxingExprIf>& core_boxing_expr) {
+  return JUST(BoxingExpr(JUST(ReplaceInDeviceType(DeviceType::kCUDA)),
+                         JUST(OptionalBoxing("copy-h2d")),
+                         JUST(BoxingExpr(JUST(ReplaceOutDeviceType(DeviceType::kCUDA)),
+                                         core_boxing_expr, JUST(OptionalBoxing("copy-d2h"))))));
+}
+
+Maybe<BoxingExprIf> SymmetricOneDimSxToBBoxingExpr() {
+  return JUST(BoxingExpr(JUST(InPlacementAndSplit(0)), JUST(OptionalBoxing("ccl-s-to-s")),
+                         JUST(BoxingExpr("ccl-s-to-b"))));
+}
+
+Maybe<BoxingExprIf> SymmetricOneDimPToSxBoxingExpr() {
+  return JUST(BoxingExpr(JUST(OutPlacementAndSplit(0)), JUST(BoxingExpr("ccl-p-to-s")),
+                         JUST(OptionalBoxing("ccl-s-to-s"))));
+}
+
+Maybe<BoxingExprIf> SymmetricCyclicNDimToNDimBoxingExpr() {
+  return JUST(BoxingExpr(JUST(InPlacementAndRepeatFirstSbp()),
+                         JUST(BoxingExpr("symmetric-acyclic-nd-sbp-to-nd-sbp")),
+                         JUST(BoxingExpr("symmetric-acyclic-nd-sbp-to-nd-sbp"))))
+         | JUST(BoxingExpr(JUST(InPlacementAndBroadcast()),
+                           JUST(BoxingExpr("symmetric-acyclic-nd-sbp-to-nd-sbp")),
+                           JUST(BoxingExpr("symmetric-acyclic-nd-sbp-to-nd-sbp"))));
+}
+
+Maybe<BoxingExprIf> SymmetricNDimToNDimBoxingExpr() {
+  return JUST(BoxingExpr("symmetric-acyclic-nd-sbp-to-nd-sbp"))
+         | JUST(SymmetricCyclicNDimToNDimBoxingExpr());
+}
+
+Maybe<BoxingExprIf> SymmetricOneDimToNDimBoxingExpr() {
+  return JUST(BoxingExpr(JUST(UnflattenInHierarchy()), JUST(BoxingExpr("unflatten-hierarchy")),
+                         JUST(SymmetricNDimToNDimBoxingExpr()) | JUST(BoxingExpr("identity"))));
+}
+
+Maybe<BoxingExprIf> SymmetricNDimToOneDimBoxingExpr() {
+  return JUST(BoxingExpr(JUST(UnflattenOutHierarchy()),
+                         JUST(SymmetricNDimToNDimBoxingExpr()) | JUST(BoxingExpr("identity")),
+                         JUST(BoxingExpr("flatten-hierarchy"))));
+}
+
+Maybe<BoxingExprIf> NToOneBoxingExpr() {
+  return JUST(BoxingExpr(JUST(InPlacementAndBroadcast()),
+                         JUST(BoxingExpr("identity")) | JUST(BoxingExpr("ccl-p-to-b"))
+                             | JUST(SymmetricOneDimSxToBBoxingExpr())
+                             | JUST(BoxingExpr("naive-p-to-b")) | JUST(BoxingExpr("naive-s-to-b"))
+                             | JUST(SymmetricNDimToNDimBoxingExpr())
+                             | JUST(BoxingExpr("generic-symmetric-nd-sbp-to-nd-sbp")),
+                         JUST(BoxingExpr("naive-b-to-1"))));
+}
+
+Maybe<BoxingExprIf> OneToNBoxingExpr() {
+  return JUST(BoxingExpr(JUST(OutPlacementAndPartialSum()), JUST(BoxingExpr("naive-1-to-p")),
+                         JUST(BoxingExpr("identity")) | JUST(BoxingExpr("ccl-p-to-b"))
+                             | JUST(SymmetricOneDimPToSxBoxingExpr())
+                             | JUST(BoxingExpr("naive-p-to-b")) | JUST(BoxingExpr("naive-p-to-s"))
+                             | JUST(SymmetricNDimToNDimBoxingExpr())
+                             | JUST(BoxingExpr("generic-symmetric-nd-sbp-to-nd-sbp"))));
+}
+
+Maybe<BoxingExprIf> SymmetricOneDimXToBBoxingExpr() {
+  return JUST(BoxingExpr("ccl-p-to-b"))
+         | JUST(BoxingExpr(JUST(InPlacementAndSplit(0)),
+                           JUST(BoxingExpr("identity")) | JUST(BoxingExpr("ccl-s-to-s")),
+                           JUST(BoxingExpr("ccl-s-to-b"))));
+}
+
+Maybe<BoxingExprIf> ASymmetricOneDimXToBBoxingExpr() {
+  return JUST(BoxingExpr(JUST(InPlacementAndBroadcast()),
+                         JUST(BoxingExpr("identity")) | JUST(SymmetricOneDimXToBBoxingExpr()),
+                         JUST(BoxingExpr("asymmetric-broadcast"))));
+}
+
+Maybe<BoxingExprIf> GenericBoxingExpr() {
+  // in_placement contain out_placement or out_placement contain in_placement
+  const auto& boxing_expr_with_inclusive_placement =
+      JUST(BoxingExpr(JUST(OutPlacementAndBroadcast()), JUST(ASymmetricOneDimXToBBoxingExpr()),
+                      JUST(BoxingExpr("identity")) | JUST(BoxingExpr("symmetric-b-to-p"))
+                          | JUST(BoxingExpr("symmetric-b-to-s"))));
+  // in_placement and out_placement have no containment relationship
+  // n to 1
+  const auto& lhs_boxing = JUST(NToOneBoxingExpr());
+  // 1 to 1 -> 1 to n
+  const auto& rhs_boxing =
+      JUST(BoxingExpr(JUST(OutFirstDeviceAndAllBroadcast()), JUST(OptionalBoxing("naive-1-to-1")),
+                      JUST(OneToNBoxingExpr())));
+  return boxing_expr_with_inclusive_placement
+         | JUST(BoxingExpr(JUST(InFirstDeviceAndAllBroadcast()), lhs_boxing, rhs_boxing));
+}
+
+Maybe<BoxingExprIf> RawMainBoxingExpr() {
+  // clang-format off
+  const auto& core = JUST(BoxingExpr("identity"))
+                     | JUST(BoxingExpr("copy-h2d"))
+                     | JUST(BoxingExpr("copy-d2h"))
+                     | JUST(BoxingExpr("ccl-p-to-b"))
+                     | JUST(BoxingExpr("ccl-s-to-s"))
+                     | JUST(SymmetricOneDimSxToBBoxingExpr())
+                     | JUST(SymmetricOneDimPToSxBoxingExpr())
+                     | JUST(BoxingExpr("symmetric-b-to-p"))
+                     | JUST(BoxingExpr("symmetric-b-to-s"))
+                     | JUST(BoxingExpr("symmetric-s-to-p"))
+                     | JUST(SymmetricOneDimXToBBoxingExpr())
+                     | JUST(ASymmetricOneDimXToBBoxingExpr())
+                     | JUST(BoxingExpr("naive-1-to-1"))
+                     | JUST(OneToNBoxingExpr())
+                     | JUST(NToOneBoxingExpr())
+                     | JUST(BoxingExpr("naive-s-to-s"))
+                     | JUST(BoxingExpr("naive-s-to-b"))
+                     | JUST(BoxingExpr("naive-b-to-s"))
+                     | JUST(BoxingExpr("naive-p-to-b"))
+                     | JUST(BoxingExpr("naive-p-to-s"))
+                     | JUST(BoxingExpr("naive-s-to-p"))
+                     | JUST(BoxingExpr("nd-sbp-dim-reduce"))
+                     | JUST(SymmetricNDimToNDimBoxingExpr())
+                     | JUST(BoxingExpr("generic-symmetric-nd-sbp-to-nd-sbp"))
+                     | JUST(SymmetricOneDimToNDimBoxingExpr())
+                     | JUST(SymmetricNDimToOneDimBoxingExpr())
+                     | JUST(GenericBoxingExpr());
+  // clang-format on
+  return core | JUST(OptionalCudaCopy(core));
+}
+
+}  // namespace
+
+static constexpr auto* MainBoxingExpr = DECORATE(&RawMainBoxingExpr, ThreadLocalCached);
+
+Maybe<EagerBoxingInterpreter> GetBoxingInterpreter(Symbol<NdSbp> in_nd_sbp,
+                                                   Symbol<NdSbp> out_nd_sbp,
+                                                   Symbol<ParallelDesc> in_parallel_desc,
+                                                   Symbol<ParallelDesc> out_parallel_desc,
+                                                   const Shape& logical_shape) {
+  const auto& in = JUST(PlacedNdSbp::New(in_nd_sbp, in_parallel_desc));
+  const auto& out = JUST(PlacedNdSbp::New(out_nd_sbp, out_parallel_desc));
+  const auto& main_boxing_expr = JUST(MainBoxingExpr());
+  const auto& status = TRY(main_boxing_expr->Check(in, out, logical_shape));
+  if (status.IsOk()) {
+    const auto& boxing_func = JUST(main_boxing_expr->GetBoxingFunction(in, out, logical_shape));
+    return std::shared_ptr<EagerBoxingInterpreter>(
+        new NaiveEagerBoxingInterpreter(boxing_func, JUST(status)));
+  }
+
+  UNIMPLEMENTED_THEN_RETURN() << Error::RuntimeError() << "global-to-global not supported"
+                              << ". from_nd_sbp: " << NdSbpToString(in_nd_sbp)
+                              << ", to_nd_sbp: " << NdSbpToString(out_nd_sbp)
+                              << ", from_placement: " << *JUST(PlacementToString(in_parallel_desc))
+                              << ", to_placement: " << *JUST(PlacementToString(out_parallel_desc));
+}
+
+static constexpr auto* CachedGetBoxingInterpreter =
+    DECORATE(&GetBoxingInterpreter, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<EagerBoxingInterpreter> EagerBoxingInterpreterManager::GetEagerBoxingInterpreter(
+    Symbol<NdSbp> in_nd_sbp, Symbol<NdSbp> out_nd_sbp, Symbol<ParallelDesc> in_parallel_desc,
+    Symbol<ParallelDesc> out_parallel_desc, const Shape& logical_shape) const {
+  return JUST(CachedGetBoxingInterpreter(in_nd_sbp, out_nd_sbp, in_parallel_desc, out_parallel_desc,
+                                         logical_shape));
+}
+
+COMMAND(
+    Singleton<EagerBoxingInterpreterManager>::SetAllocated(new EagerBoxingInterpreterManager()));
+
+}  // namespace oneflow

oneflow/core/boxing/eager_boxing_interpreter_mgr.h

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_BOXING_EAGER_BOXING_INTERPRETER_MGR_H_
+#define ONEFLOW_CORE_BOXING_EAGER_BOXING_INTERPRETER_MGR_H_
+
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+
+namespace oneflow {
+
+class EagerBoxingInterpreterManager final {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(EagerBoxingInterpreterManager);
+  EagerBoxingInterpreterManager() = default;
+  virtual ~EagerBoxingInterpreterManager() = default;
+
+  Maybe<EagerBoxingInterpreter> GetEagerBoxingInterpreter(Symbol<NdSbp> in_nd_sbp,
+                                                          Symbol<NdSbp> out_nd_sbp,
+                                                          Symbol<ParallelDesc> in_parallel_desc,
+                                                          Symbol<ParallelDesc> out_parallel_desc,
+                                                          const Shape& logical_shape) const;
+};
+
+template<typename RetT, typename... Args>
+struct DisableRecusiveBoxingCall {
+  static_assert(is_maybe<RetT>::value, "returned value type must be Maybe<T>.");
+  template<RetT (*func)(Args...)>
+  static RetT Call(Args... arg) {
+    static thread_local bool disable_boxing = false;
+    CHECK_OR_RETURN(!disable_boxing);
+    disable_boxing = true;
+    RetT ret = func(arg...);
+    disable_boxing = false;
+    return ret;
+  }
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_BOXING_EAGER_BOXING_INTERPRETER_MGR_H_

oneflow/core/boxing/eager_boxing_logger.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/singleton.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/common/env_var/debug_mode.h"
+#include "oneflow/core/boxing/eager_boxing_logger.h"
+#include "oneflow/core/boxing/boxing_interpreter_status.h"
+
+namespace oneflow {
+
+namespace {
+
+class NullEagerBoxingLogger final : public EagerBoxingLogger {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(NullEagerBoxingLogger);
+  NullEagerBoxingLogger() = default;
+  ~NullEagerBoxingLogger() override = default;
+
+  void Log(const BoxingInterpreterStatus& status, const std::string& prefix) const override {}
+};
+
+class NaiveEagerBoxingLogger final : public EagerBoxingLogger {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(NaiveEagerBoxingLogger);
+  NaiveEagerBoxingLogger() = default;
+  ~NaiveEagerBoxingLogger() override = default;
+
+  void Log(const BoxingInterpreterStatus& status, const std::string& prefix) const override {
+    LOG(INFO) << prefix << "Boxing route: " << (status.boxing_routing());
+    LOG(INFO) << prefix << "Logical shape: " << (status.logical_shape().ToString());
+    LOG(INFO) << prefix << "Altered state of sbp: " << (status.nd_sbp_routing());
+    LOG(INFO) << prefix << "Altered state of placement: " << (status.placement_routing());
+  }
+};
+
+const EagerBoxingLogger* CreateEagerBoxingLogger() {
+  if (IsInDebugMode()) {
+    return new NaiveEagerBoxingLogger();
+  } else {
+    return new NullEagerBoxingLogger();
+  }
+}
+
+}  // namespace
+
+COMMAND(Singleton<const EagerBoxingLogger>::SetAllocated(CreateEagerBoxingLogger()));
+
+}  // namespace oneflow

oneflow/core/boxing/eager_boxing_logger.h

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_BOXING_EAGER_BOXING_LOGGER_H_
+#define ONEFLOW_CORE_BOXING_EAGER_BOXING_LOGGER_H_
+
+#include "oneflow/core/common/util.h"
+
+namespace oneflow {
+
+class BoxingInterpreterStatus;
+
+class EagerBoxingLogger {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(EagerBoxingLogger);
+  EagerBoxingLogger() = default;
+  virtual ~EagerBoxingLogger() = default;
+
+  virtual void Log(const BoxingInterpreterStatus& status, const std::string& prefix) const = 0;
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_BOXING_EAGER_BOXING_LOGGER_H_

oneflow/core/boxing/flatten_hierarchy.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/operator/operator.h"
+
+namespace oneflow {
+
+namespace {
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckFlattenHierarchy(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                     const Shape& logical_shape) {
+  CHECK_GT_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  for (int i = 0; i < in->nd_sbp()->sbp_parallel_size(); ++i) {
+    const auto& sbp_parallel = in->nd_sbp()->sbp_parallel(i);
+    CHECK_OR_RETURN(sbp_parallel == out->nd_sbp()->sbp_parallel(0)) << "nd_sbp axis: " << i;
+  }
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), out->placement()->device_type());
+  CHECK_EQ_OR_RETURN(in->placement()->parallel_num(), out->placement()->parallel_num());
+  ParallelConf flattened_parallel_conf(in->placement()->parallel_conf());
+  flattened_parallel_conf.clear_hierarchy();
+  const auto& flatten_placement = SymbolOf(ParallelDesc(flattened_parallel_conf));
+  CHECK_OR_RETURN(flatten_placement == out->placement())
+      << "The output placement is not a hierarch-flattened version of the input placement";
+  for (int64_t in_parallel_id = 0; in_parallel_id < in->placement()->parallel_num();
+       ++in_parallel_id) {
+    const auto& in_physical_shape =
+        JUST(GetPhysicalShape(logical_shape, *in->nd_sbp(), *in->placement(), in_parallel_id));
+    const auto& out_physical_shape =
+        JUST(GetPhysicalShape(logical_shape, *out->nd_sbp(), *out->placement(), in_parallel_id));
+    CHECK_EQ_OR_RETURN(*in_physical_shape, *out_physical_shape);
+  }
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+}  // namespace
+
+static constexpr auto* CheckFlattenHierarchy =
+    DECORATE(&RawCheckFlattenHierarchy, ThreadLocalCachedCopiable);
+
+Maybe<one::Tensor> FlattenHierarchy(const std::shared_ptr<one::Tensor>& tensor,
+                                    Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  const auto& local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+COMMAND(RegisterBoxingFunction("flatten-hierarchy", CheckFlattenHierarchy, &FlattenHierarchy));
+
+}  // namespace oneflow

oneflow/core/boxing/generic_symmetric_nd_sbp_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/boxing/eager_boxing_interpreter_mgr.h"
+#include "oneflow/core/boxing/eager_boxing_logger.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/framework/placement_sbp_util.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/operator/operator.h"
+#include "oneflow/core/common/stride.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsAllBroadcastNdSbpAfterDim(Symbol<NdSbp> nd_sbp, int dim) {
+  for (int i = dim; i < nd_sbp->sbp_parallel_size(); ++i) {
+    if (!nd_sbp->sbp_parallel(i).has_broadcast_parallel()) { return false; }
+  }
+  return true;
+}
+
+static constexpr auto* IsAllBroadcastNdSbpAfterDim =
+    DECORATE(&RawIsAllBroadcastNdSbpAfterDim, ThreadLocalCached);
+
+Maybe<Symbol<SbpParallel>> GetBroadcastSbp() {
+  SbpParallel broadcast_sbp;
+  broadcast_sbp.mutable_broadcast_parallel();
+  return SymbolOf(broadcast_sbp);
+}
+
+auto* CachedGetBroadcastSbp = DECORATE(&GetBroadcastSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<Shape> CalcLogicalShape4Axis(const Shape& logical_shape, int axis,
+                                   Symbol<ParallelDesc> parallel_desc, Symbol<NdSbp> nd_sbp) {
+  CHECK_LT_OR_RETURN(axis, nd_sbp->sbp_parallel_size());  // Always true
+  std::shared_ptr<Shape> sub_logical_shape = std::make_shared<Shape>(logical_shape);
+
+  const auto& opt_parallel_id = JUST(GetParallelId4CurrentProcessCtx(parallel_desc));
+  int64_t parallel_id = JUST(*opt_parallel_id);
+  const auto& hierarchy_shape = *parallel_desc->hierarchy();
+  Stride hierarchy_stride(hierarchy_shape);
+
+  FOR_RANGE(int64_t, i, 0, axis) {
+    const auto& sbp_parallel = nd_sbp->sbp_parallel(i);
+    if (sbp_parallel.has_split_parallel()) {
+      int64_t index = CalcIndex4Axis(parallel_id, hierarchy_stride, i);
+      int64_t dim = hierarchy_shape.At(i);
+      const int64_t split_axis = sbp_parallel.split_parallel().axis();
+
+      if (sub_logical_shape->At(split_axis) > 0) {
+        CHECK_GE_OR_RETURN(sub_logical_shape->At(split_axis), dim)
+            << Error::RuntimeError() << "The size of tensor (" << sub_logical_shape->At(split_axis)
+            << ") at split dimension (" << i
+            << ") should be greater than or equal to parallle num (" << dim << ")";
+        const BalancedSplitter bs(sub_logical_shape->At(split_axis), dim);
+        sub_logical_shape->Set(split_axis, bs.At(index).size());
+      }
+    }
+  }
+
+  return sub_logical_shape;
+}
+
+static constexpr auto* GetLogicalShape4Axis =
+    DECORATE(&CalcLogicalShape4Axis, ThreadLocalCachedCopiable);
+
+Maybe<int> CalcTheFirstDiffAxisBetweenTwoNdSbp(Symbol<NdSbp> in_nd_sbp, Symbol<NdSbp> out_nd_sbp) {
+  CHECK_EQ_OR_RETURN(in_nd_sbp->sbp_parallel_size(),
+                     out_nd_sbp->sbp_parallel_size());  // Always true
+  int dim = 0;
+  for (; dim < in_nd_sbp->sbp_parallel_size(); ++dim) {
+    if (in_nd_sbp->sbp_parallel(dim) != out_nd_sbp->sbp_parallel(dim)) { break; }
+  }
+  return dim;
+}
+
+Maybe<one::Tensor> Apply1DBoxing(const std::shared_ptr<one::Tensor>& input, Symbol<NdSbp> in_nd_sbp,
+                                 Symbol<NdSbp> out_nd_sbp, Symbol<ParallelDesc> in_parallel_desc,
+                                 Symbol<ParallelDesc> out_parallel_desc) {
+  const auto& boxing_interpreter =
+      JUST(Singleton<EagerBoxingInterpreterManager>::Get()->GetEagerBoxingInterpreter(
+          in_nd_sbp, out_nd_sbp, in_parallel_desc, out_parallel_desc, *input->shape()));
+  Singleton<const EagerBoxingLogger>::Get()->Log(
+      *JUST(boxing_interpreter->boxing_interpreter_status()),
+      /* prefix */ "\t\tInternal boxing of generic-symmetric-nd-sbp-to-nd-sbp, ");
+  return JUST(boxing_interpreter->Interpret(input, in_nd_sbp, out_nd_sbp, in_parallel_desc,
+                                            out_parallel_desc));
+}
+
+Maybe<void> RawCheckGenericSymmetricNdSbpBoxing(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                                const Shape& logical_shape) {
+  CHECK_OR_RETURN(in->placement() == out->placement());
+  CHECK_OR_RETURN(in->nd_sbp() != out->nd_sbp());
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), out->nd_sbp()->sbp_parallel_size());
+  CHECK_GT_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckGenericSymmetricNdSbpBoxing =
+    DECORATE(&RawCheckGenericSymmetricNdSbpBoxing, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> GenericSymmetricNdSbpBoxing(const std::shared_ptr<one::Tensor>& input,
+                                               Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) {
+  const auto& in_parallel_desc = in->placement();
+  const auto& out_nd_sbp = out->nd_sbp();
+  const auto& out_parallel_desc = out->placement();
+  std::shared_ptr<one::Tensor> output;
+
+  const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out_parallel_desc));
+  if (out_parallel_id->has_value()) {
+    output = input;
+
+    int first_diff_sbp_dim = JUST(CalcTheFirstDiffAxisBetweenTwoNdSbp(in->nd_sbp(), out_nd_sbp));
+    Symbol<SbpParallel> broadcast_sbp = JUST(CachedGetBroadcastSbp());
+
+    const auto& opt_parallel_id = JUST(GetParallelId4CurrentProcessCtx(in_parallel_desc));
+    int64_t parallel_id = JUST(*opt_parallel_id);
+    const auto& hierarchy_shape = *in_parallel_desc->hierarchy();
+    Stride hierarchy_stride(hierarchy_shape);
+
+    const auto& logical_shape = input->shape();
+
+    // Convert input to broadcast tensor step by step
+    // e.g.
+    // If in_nd_sbp is (S(0), B, S(0)), (S(0), S(0), S(1))
+    // Altered state of sbp is (S(0), B, S(0)) -> (S(0), B, B)
+    for (int64_t i = out_nd_sbp->sbp_parallel_size() - 1; i >= first_diff_sbp_dim; --i) {
+      const auto& nd_sbp = JUST(output->nd_sbp());
+      const auto& sbp_parallel = nd_sbp->sbp_parallel(i);
+      if (sbp_parallel.has_broadcast_parallel()) { continue; }
+
+      const auto& one_dim_nd_sbp = JUST(SbpToNdSbp(sbp_parallel));
+      const auto& sub_logical_shape =
+          *JUST(GetLogicalShape4Axis(*logical_shape, i, in_parallel_desc, nd_sbp));
+      std::shared_ptr<one::Tensor> local_tensor = JUST(output->cur_rank_phy_tensor());
+      const auto& sub_parallel_desc = JUST(CalcSubParallelDesc4Axis(in_parallel_desc, i));
+
+      int64_t index = CalcIndex4Axis(parallel_id, hierarchy_stride, i);
+
+      const auto& physical_shape =
+          JUST(GetPhysicalShape(sub_logical_shape, *one_dim_nd_sbp, *sub_parallel_desc, index));
+      CHECK_EQ_OR_RETURN(*physical_shape, *local_tensor->shape())
+          << Error::RuntimeError() << "Invalid input tensor, size of local tensor ("
+          << local_tensor->shape()->ToString() << ") does not match global tensor ("
+          << logical_shape->ToString() << ")!";
+      std::shared_ptr<one::Tensor> sub_global_tensor = JUST(one::functional::LocalToConsistent(
+          local_tensor, sub_parallel_desc, *JUST(GetSbpList(one_dim_nd_sbp)), sub_logical_shape,
+          local_tensor->dtype()));
+
+      sub_global_tensor =
+          JUST(Apply1DBoxing(sub_global_tensor, one_dim_nd_sbp, JUST(SbpToNdSbp(broadcast_sbp)),
+                             sub_parallel_desc, sub_parallel_desc));
+
+      local_tensor = JUST(sub_global_tensor->cur_rank_phy_tensor());
+
+      const auto& new_nd_sbp = JUST(SetSbpAtAxis(*nd_sbp, *broadcast_sbp, i));
+
+      output = JUST(one::functional::LocalToConsistent(local_tensor, in_parallel_desc,
+                                                       *JUST(GetSbpList(new_nd_sbp)),
+                                                       *logical_shape, local_tensor->dtype()));
+    }
+
+    CHECK_OR_RETURN(IsAllBroadcastNdSbpAfterDim(JUST(output->nd_sbp()), first_diff_sbp_dim))
+        << Error::RuntimeError()
+        << "Compute generic-symmetric-nd-sbp-to-nd-sbp failed. Please submit an issue in "
+           "`https://github.com/Oneflow-Inc/oneflow/issues` and we will fix it as soon as "
+           "possible";
+
+    // Convert broadcast tensor to output with out_nd_sbp data step by step
+    // e.g.
+    // If out_nd_sbp is (S(0), S(0), S(1))
+    // Altered state of sbp is (S(0), B, B) -> (S(0), S(0), B) -> (S(0), S(0), S(1))
+    std::shared_ptr<Shape> sub_logical_shape = JUST(GetLogicalShape4Axis(
+        *logical_shape, first_diff_sbp_dim, in_parallel_desc, JUST(output->nd_sbp())));
+    for (int64_t i = first_diff_sbp_dim; i < out_nd_sbp->sbp_parallel_size(); ++i) {
+      const auto& sbp_parallel = out_nd_sbp->sbp_parallel(i);
+      if (sbp_parallel.has_broadcast_parallel()) { continue; }
+
+      const auto& nd_sbp = JUST(output->nd_sbp());
+
+      const auto& sub_parallel_desc = JUST(CalcSubParallelDesc4Axis(in_parallel_desc, i));
+
+      std::shared_ptr<one::Tensor> local_tensor = JUST(output->cur_rank_phy_tensor());
+
+      std::shared_ptr<one::Tensor> sub_global_tensor = JUST(one::functional::LocalToConsistent(
+          local_tensor, sub_parallel_desc, *JUST(GetSbpList(JUST(SbpToNdSbp(broadcast_sbp)))),
+          *sub_logical_shape, local_tensor->dtype()));
+
+      const auto& one_dim_nd_sbp = JUST(SbpToNdSbp(sbp_parallel));
+      sub_global_tensor = JUST(Apply1DBoxing(sub_global_tensor, JUST(SbpToNdSbp(broadcast_sbp)),
+                                             one_dim_nd_sbp, sub_parallel_desc, sub_parallel_desc));
+
+      local_tensor = JUST(sub_global_tensor->cur_rank_phy_tensor());
+
+      int64_t index = CalcIndex4Axis(parallel_id, hierarchy_stride, i);
+      const auto& physical_shape =
+          JUST(GetPhysicalShape(*sub_logical_shape, *one_dim_nd_sbp, *sub_parallel_desc, index));
+      CHECK_EQ_OR_RETURN(*physical_shape, *local_tensor->shape())
+          << Error::RuntimeError()
+          << "Compute generic-symmetric-nd-sbp-to-nd-sbp failed. Please submit an issue in "
+             "`https://github.com/Oneflow-Inc/oneflow/issues` and we will fix it as soon as "
+             "possible";
+
+      const auto& new_nd_sbp = JUST(SetSbpAtAxis(*nd_sbp, sbp_parallel, i));
+
+      output = JUST(one::functional::LocalToConsistent(local_tensor, in_parallel_desc,
+                                                       *JUST(GetSbpList(new_nd_sbp)),
+                                                       *logical_shape, local_tensor->dtype()));
+      // physical_shape of this axis is logical shape of next axis
+      sub_logical_shape = physical_shape;
+    }
+  } else {
+    one::ConsistentTensorMeta tensor_meta(input->shape(), input->dtype()->data_type(), out_nd_sbp,
+                                          out_parallel_desc);
+    const auto& tensor_impl = JUST(
+        one::EagerConsistentTensorImpl::New(SymbolOf(tensor_meta), input->requires_grad(), false));
+    output = std::make_shared<one::ConsistentTensor>(tensor_impl);
+  }
+
+  return output;
+}
+
+COMMAND(RegisterBoxingFunction("generic-symmetric-nd-sbp-to-nd-sbp",
+                               CheckGenericSymmetricNdSbpBoxing, &GenericSymmetricNdSbpBoxing));
+
+}  // namespace oneflow

oneflow/core/boxing/identity_boxing_interpreter.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/framework/nd_sbp.h"
+namespace oneflow {
+
+namespace {
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckIdentity(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                             const Shape& logical_shape) {
+  if (in->placement()->parallel_num() == 1) {
+    CHECK_OR_RETURN(in->placement()->EqualsIgnoringHierarchy(*out->placement()));
+    return Maybe<void>::Ok();
+  }
+  CHECK_OR_RETURN(in->placement() == out->placement());
+  CHECK_OR_RETURN(in->nd_sbp() == out->nd_sbp());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+}  // namespace
+
+Maybe<one::Tensor> GetIdentity(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                               Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  // reset sbp if parallel_num == 1 and reset transport_token
+  const auto& local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+COMMAND(RegisterBoxingFunction("identity", DECORATE(&RawCheckIdentity, ThreadLocalCachedCopiable),
+                               &GetIdentity));
+}  // namespace oneflow

oneflow/core/boxing/naive_1_to_p_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/control/global_process_ctx.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/framework/device.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/common/decorator.h"
+
+namespace oneflow {
+
+namespace {
+
+bool NdSbpIsAllPartialSum(Symbol<NdSbp> nd_sbp) {
+  for (const auto& sbp_parallel : nd_sbp->sbp_parallel()) {
+    if (!sbp_parallel.has_partial_sum_parallel()) { return false; }
+  }
+  return true;
+}
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaive1ToP(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->placement()->parallel_num(), 1);
+  CHECK_OR_RETURN(NdSbpIsAllPartialSum(out->nd_sbp()));
+  CHECK_OR_RETURN(out->placement()->Bigger(*in->placement()));
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaive1ToP = DECORATE(&RawCheckNaive1ToP, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> Naive1ToP(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+
+  int64_t root = JUST(tensor_placement->MachineId4ParallelId(0));
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+  if (root == GlobalProcessCtx::Rank() || !out_parallel_id->has_value()) {
+    // do nothing
+  } else {
+    const std::string& device_type = tensor_placement->device_tag();
+    local_tensor = JUST(one::functional::Constant(*tensor->shape(), 0, tensor->dtype(),
+                                                  JUST(Device::New(device_type))));
+  }
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(),
+                                                 *JUST(GetSbpList(out->nd_sbp())), *tensor->shape(),
+                                                 tensor->dtype()));
+}
+
+COMMAND(RegisterBoxingFunction("naive-1-to-p", CheckNaive1ToP, &Naive1ToP));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_b_to_1_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/control/global_process_ctx.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/job/nd_sbp_util.h"
+#include "oneflow/core/framework/device.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/common/decorator.h"
+
+namespace oneflow {
+
+namespace {
+
+Maybe<void> RawCheckNaiveBTo1(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  // NOLINTBEGIN(maybe-need-error-msg)
+  CHECK_EQ_OR_RETURN(out->placement()->parallel_num(), 1);
+  CHECK_OR_RETURN(NdSbpIsAllBroadcast(*in->nd_sbp()));
+  CHECK_OR_RETURN(in->placement()->Bigger(*out->placement()));
+  // NOLINTEND(maybe-need-error-msg)
+  return Maybe<void>::Ok();
+}
+
+static constexpr auto* CheckNaiveBTo1 = DECORATE(&RawCheckNaiveBTo1, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaiveBTo1(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(),
+                                                 *JUST(GetSbpList(out->nd_sbp())), *tensor->shape(),
+                                                 tensor->dtype()));
+}
+
+COMMAND(RegisterBoxingFunction("naive-b-to-1", CheckNaiveBTo1, &NaiveBTo1));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_b_to_s_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/boxing/slice_boxing_util.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsSplitSbp(Symbol<SbpParallel> sbp_parallel) { return sbp_parallel->has_split_parallel(); }
+
+static constexpr auto* IsSplitSbp = DECORATE(&RawIsSplitSbp, ThreadLocalCached);
+
+bool RawIsBroadcastSbp(Symbol<SbpParallel> sbp_parallel) {
+  return sbp_parallel->has_broadcast_parallel();
+}
+
+static constexpr auto* IsBroadcastSbp = DECORATE(&RawIsBroadcastSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaiveBToS(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+
+  CHECK_OR_RETURN(IsBroadcastSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsSplitSbp(out->nd_sbp()->sbp_parallel(0)));
+
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaiveBToS = DECORATE(&RawCheckNaiveBToS, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaiveBToS(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  {
+    const auto& in_parallel_id = JUST(GetParallelId4CurrentProcessCtx(tensor_placement));
+    const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+    if (in_parallel_id->has_value() || out_parallel_id->has_value()) {
+      local_tensor = JUST(one::functional::EagerBToS(
+          local_tensor, tensor_placement, out->placement(), *sbp_list, *tensor->shape()));
+    }
+  }
+
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+static constexpr auto* NaiveBToSWithAutoConvert =
+    EAGER_SLICE_BOXING_WARPPER(&NaiveBToS, EagerSliceBoxingType::kNaiveBToS);
+
+COMMAND(RegisterBoxingFunction("naive-b-to-s", CheckNaiveBToS, NaiveBToSWithAutoConvert));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_p_to_b_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/boxing/slice_boxing_util.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsPartialSumSbp(Symbol<SbpParallel> sbp_parallel) {
+  return sbp_parallel->has_partial_sum_parallel();
+}
+
+static constexpr auto* IsPartialSumSbp = DECORATE(&RawIsPartialSumSbp, ThreadLocalCached);
+
+bool RawIsBroadcastSbp(Symbol<SbpParallel> sbp_parallel) {
+  return sbp_parallel->has_broadcast_parallel();
+}
+
+static constexpr auto* IsBroadcastSbp = DECORATE(&RawIsBroadcastSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaivePToB(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_OR_RETURN(IsPartialSumSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsBroadcastSbp(out->nd_sbp()->sbp_parallel(0)));
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaivePToB = DECORATE(&RawCheckNaivePToB, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaivePToB(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  {
+    const auto& in_parallel_id = JUST(GetParallelId4CurrentProcessCtx(tensor_placement));
+    const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+    if (in_parallel_id->has_value() || out_parallel_id->has_value()) {
+      local_tensor = JUST(one::functional::EagerPToB(local_tensor, tensor_placement,
+                                                     out->placement(), *tensor->shape()));
+    }
+  }
+
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+static constexpr auto* NaivePToBWithAutoConvert =
+    EAGER_SLICE_BOXING_WARPPER(&NaivePToB, EagerSliceBoxingType::kNaivePToB);
+
+COMMAND(RegisterBoxingFunction("naive-p-to-b", CheckNaivePToB, NaivePToBWithAutoConvert));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_p_to_s_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/boxing/slice_boxing_util.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsPartialSumSbp(Symbol<SbpParallel> sbp_parallel) {
+  return sbp_parallel->has_partial_sum_parallel();
+}
+
+static constexpr auto* IsPartialSumSbp = DECORATE(&RawIsPartialSumSbp, ThreadLocalCached);
+
+bool RawIsSplitSbp(Symbol<SbpParallel> sbp_parallel) { return sbp_parallel->has_split_parallel(); }
+
+static constexpr auto* IsSplitSbp = DECORATE(&RawIsSplitSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaivePToS(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+
+  CHECK_OR_RETURN(IsPartialSumSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsSplitSbp(out->nd_sbp()->sbp_parallel(0)));
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaivePToS = DECORATE(&RawCheckNaivePToS, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaivePToS(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  {
+    const auto& in_parallel_id = JUST(GetParallelId4CurrentProcessCtx(tensor_placement));
+    const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+    if (in_parallel_id->has_value() || out_parallel_id->has_value()) {
+      local_tensor = JUST(one::functional::EagerPToS(
+          local_tensor, tensor_placement, out->placement(), *sbp_list, *tensor->shape()));
+    }
+  }
+
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+static constexpr auto* NaivePToSWithAutoConvert =
+    EAGER_SLICE_BOXING_WARPPER(&NaivePToS, EagerSliceBoxingType::kNaivePToS);
+
+COMMAND(RegisterBoxingFunction("naive-p-to-s", CheckNaivePToS, NaivePToSWithAutoConvert));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_s_to_b_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/boxing/slice_boxing_util.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsSplitSbp(Symbol<SbpParallel> sbp_parallel) { return sbp_parallel->has_split_parallel(); }
+
+static constexpr auto* IsSplitSbp = DECORATE(&RawIsSplitSbp, ThreadLocalCached);
+
+bool RawIsBroadcastSbp(Symbol<SbpParallel> sbp_parallel) {
+  return sbp_parallel->has_broadcast_parallel();
+}
+
+static constexpr auto* IsBroadcastSbp = DECORATE(&RawIsBroadcastSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaiveSToB(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_OR_RETURN(IsSplitSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsBroadcastSbp(out->nd_sbp()->sbp_parallel(0)));
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaiveSToB = DECORATE(&RawCheckNaiveSToB, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaiveSToB(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  {
+    const auto& in_parallel_id = JUST(GetParallelId4CurrentProcessCtx(tensor_placement));
+    const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+    if (in_parallel_id->has_value() || out_parallel_id->has_value()) {
+      local_tensor =
+          JUST(one::functional::EagerSToB(local_tensor, tensor_placement, out->placement(),
+                                          *JUST(GetSbpList(tensor_nd_sbp)), *tensor->shape()));
+    }
+  }
+
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+static constexpr auto* NaiveSToBWithAutoConvert =
+    EAGER_SLICE_BOXING_WARPPER(&NaiveSToB, EagerSliceBoxingType::kNaiveSToB);
+
+COMMAND(RegisterBoxingFunction("naive-s-to-b", CheckNaiveSToB, NaiveSToBWithAutoConvert));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_s_to_p_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/boxing/slice_boxing_util.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsSplitSbp(Symbol<SbpParallel> sbp_parallel) { return sbp_parallel->has_split_parallel(); }
+
+static constexpr auto* IsSplitSbp = DECORATE(&RawIsSplitSbp, ThreadLocalCached);
+
+bool RawIsPartialSumSbp(Symbol<SbpParallel> sbp_parallel) {
+  return sbp_parallel->has_partial_sum_parallel();
+}
+
+static constexpr auto* IsPartialSumSbp = DECORATE(&RawIsPartialSumSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaiveSToP(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_OR_RETURN(IsSplitSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsPartialSumSbp(out->nd_sbp()->sbp_parallel(0)));
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaiveSToP = DECORATE(&RawCheckNaiveSToP, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaiveSToP(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  {
+    const auto& in_parallel_id = JUST(GetParallelId4CurrentProcessCtx(tensor_placement));
+    const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+    if (in_parallel_id->has_value() || out_parallel_id->has_value()) {
+      local_tensor =
+          JUST(one::functional::EagerSToP(local_tensor, tensor_placement, out->placement(),
+                                          *JUST(GetSbpList(tensor_nd_sbp)), *tensor->shape()));
+    }
+  }
+
+  const auto& sbp_list = JUST(GetSbpList(out->nd_sbp()));
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+static constexpr auto* NaiveSToPWithAutoConvert =
+    EAGER_SLICE_BOXING_WARPPER(&NaiveSToP, EagerSliceBoxingType::kNaiveSToP);
+
+COMMAND(RegisterBoxingFunction("naive-s-to-p", CheckNaiveSToP, NaiveSToPWithAutoConvert));
+
+}  // namespace oneflow

oneflow/core/boxing/naive_s_to_s_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/balanced_splitter.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/boxing/slice_boxing_util.h"
+
+namespace oneflow {
+
+namespace {
+
+bool RawIsSplitSbp(Symbol<SbpParallel> sbp_parallel) { return sbp_parallel->has_split_parallel(); }
+
+static constexpr auto* IsSplitSbp = DECORATE(&RawIsSplitSbp, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaiveSToS(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                              const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+
+  CHECK_OR_RETURN(IsSplitSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsSplitSbp(out->nd_sbp()->sbp_parallel(0)));
+
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaiveSToS = DECORATE(&RawCheckNaiveSToS, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaiveSToS(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                             Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+  const auto& in_sbp_list = JUST(GetSbpList(tensor_nd_sbp));
+  const auto& out_sbp_list = JUST(GetSbpList(out->nd_sbp()));
+
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  {
+    const auto& in_parallel_id = JUST(GetParallelId4CurrentProcessCtx(tensor_placement));
+    const auto& out_parallel_id = JUST(GetParallelId4CurrentProcessCtx(out->placement()));
+    if (in_parallel_id->has_value() || out_parallel_id->has_value()) {
+      local_tensor =
+          JUST(one::functional::EagerNaiveSToS(local_tensor, tensor_placement, out->placement(),
+                                               *in_sbp_list, *out_sbp_list, *tensor->shape()));
+    }
+  }
+
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(), *out_sbp_list,
+                                                 *tensor->shape(), tensor->dtype()));
+}
+
+static constexpr auto* NaiveSToSWithAutoConvert =
+    EAGER_SLICE_BOXING_WARPPER(&NaiveSToS, EagerSliceBoxingType::kNaiveSToS);
+
+COMMAND(RegisterBoxingFunction("naive-s-to-s", CheckNaiveSToS, NaiveSToSWithAutoConvert));
+
+}  // namespace oneflow

oneflow/core/boxing/nccl_boxing_function.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/job/nd_sbp_util.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+
+namespace {
+
+bool IsSplitSbp(Symbol<SbpParallel> sbp_parallel) { return sbp_parallel->has_split_parallel(); }
+
+Maybe<void> RawCheckNcclP2B(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                            const Shape& logical_shape) {
+  // NOLINTBEGIN(maybe-need-error-msg)
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_OR_RETURN(NdSbpIsAllPartialSum(*in->nd_sbp()));
+  CHECK_OR_RETURN(NdSbpIsAllBroadcast(*out->nd_sbp()));
+
+  CHECK_OR_RETURN(in->placement() == out->placement());
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), DeviceType::kCUDA);
+  // NOLINTEND(maybe-need-error-msg)
+  return Maybe<void>::Ok();
+}
+
+static constexpr auto* CheckNcclP2B = DECORATE(&RawCheckNcclP2B, ThreadLocalCachedCopiable);
+
+Maybe<void> RawCheckNcclP2S(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                            const Shape& logical_shape) {
+  // NOLINTBEGIN(maybe-need-error-msg)
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_OR_RETURN(NdSbpIsAllPartialSum(*in->nd_sbp()));
+  CHECK_OR_RETURN(NdSbpIsAllSplit(*out->nd_sbp(), 0));
+
+  CHECK_GT_OR_RETURN(logical_shape.NumAxes(), 0);
+  CHECK_OR_RETURN(logical_shape.At(0) % in->placement()->parallel_num() == 0);
+
+  CHECK_OR_RETURN(in->placement() == out->placement());
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), DeviceType::kCUDA);
+  // NOLINTEND(maybe-need-error-msg)
+  return Maybe<void>::Ok();
+}
+
+static constexpr auto* CheckNcclP2S = DECORATE(&RawCheckNcclP2S, ThreadLocalCachedCopiable);
+
+Maybe<void> RawCheckNcclS2B(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                            const Shape& logical_shape) {
+  // NOLINTBEGIN(maybe-need-error-msg)
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_OR_RETURN(NdSbpIsAllSplit(*in->nd_sbp(), 0));
+  CHECK_OR_RETURN(NdSbpIsAllBroadcast(*out->nd_sbp()));
+
+  CHECK_GT_OR_RETURN(logical_shape.NumAxes(), 0);
+  CHECK_OR_RETURN(logical_shape.At(0) % in->placement()->parallel_num() == 0);
+
+  CHECK_OR_RETURN(in->placement() == out->placement());
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), DeviceType::kCUDA);
+  // NOLINTEND(maybe-need-error-msg)
+  return Maybe<void>::Ok();
+}
+
+static constexpr auto* CheckNcclS2B = DECORATE(&RawCheckNcclS2B, ThreadLocalCachedCopiable);
+
+Maybe<void> RawCheckNcclS2S(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                            const Shape& logical_shape) {
+  // NOLINTBEGIN(maybe-need-error-msg)
+  CHECK_EQ_OR_RETURN(in->nd_sbp()->sbp_parallel_size(), 1);
+  CHECK_EQ_OR_RETURN(out->nd_sbp()->sbp_parallel_size(), 1);
+
+  CHECK_OR_RETURN(IsSplitSbp(in->nd_sbp()->sbp_parallel(0)));
+  CHECK_OR_RETURN(IsSplitSbp(out->nd_sbp()->sbp_parallel(0)));
+  CHECK_NE_OR_RETURN(in->nd_sbp()->sbp_parallel(0).split_parallel().axis(),
+                     out->nd_sbp()->sbp_parallel(0).split_parallel().axis());
+
+  int64_t in_split_axis = in->nd_sbp()->sbp_parallel(0).split_parallel().axis();
+  int64_t out_split_axis = out->nd_sbp()->sbp_parallel(0).split_parallel().axis();
+  CHECK_GT_OR_RETURN(logical_shape.NumAxes(), in_split_axis);
+  CHECK_GT_OR_RETURN(logical_shape.NumAxes(), out_split_axis);
+  CHECK_OR_RETURN(logical_shape.At(in_split_axis) % in->placement()->parallel_num() == 0);
+  CHECK_OR_RETURN(logical_shape.At(out_split_axis) % in->placement()->parallel_num() == 0);
+
+  CHECK_OR_RETURN(in->placement() == out->placement());
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), DeviceType::kCUDA);
+  // NOLINTEND(maybe-need-error-msg)
+  return Maybe<void>::Ok();
+}
+
+static constexpr auto* CheckNcclS2S = DECORATE(&RawCheckNcclS2S, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NcclP2B(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                           Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp());  // NOLINT(maybe-need-error-msg)
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement());  // NOLINT(maybe-need-error-msg)
+
+  return JUST(one::functional::ConsistentAllReduce(tensor));
+}
+
+Maybe<one::Tensor> NcclP2S(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                           Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp());  // NOLINT(maybe-need-error-msg)
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement());  // NOLINT(maybe-need-error-msg)
+
+  return JUST(one::functional::ConsistentReduceScatter(tensor, "sum"));
+}
+
+Maybe<one::Tensor> NcclS2B(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                           Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp());  // NOLINT(maybe-need-error-msg)
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement());  // NOLINT(maybe-need-error-msg)
+
+  return JUST(one::functional::ConsistentAllGather(tensor));
+}
+
+Maybe<one::Tensor> NcclS2S(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                           Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp());  // NOLINT(maybe-need-error-msg)
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement());  // NOLINT(maybe-need-error-msg)
+  return JUST(one::functional::ConsistentS2S(tensor, *JUST(GetSbpList(out->nd_sbp()))));
+}
+
+COMMAND(RegisterBoxingFunction("nccl-p-to-b", CheckNcclP2B, &NcclP2B));
+COMMAND(RegisterBoxingFunction("nccl-p-to-s", CheckNcclP2S, &NcclP2S));
+COMMAND(RegisterBoxingFunction("nccl-s-to-b", CheckNcclS2B, &NcclS2B));
+COMMAND(RegisterBoxingFunction("nccl-s-to-s", CheckNcclS2S, &NcclS2S));
+
+}  // namespace oneflow

oneflow/core/boxing/nd_sbp_dim_reduce_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/boxing/eager_boxing_interpreter_mgr.h"
+#include "oneflow/core/boxing/eager_boxing_logger.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/framework/device.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/operator/operator.h"
+#include "oneflow/core/framework/sbp_infer_util.h"
+
+namespace oneflow {
+
+namespace {
+
+Maybe<std::tuple<Symbol<PlacedNdSbp>, Symbol<PlacedNdSbp>>> RawInOutPlacedNdSbpDimReduce(
+    Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) {
+  // reduce hierarchy
+  ParallelDesc reduced_in_placement = *in->placement();
+  ParallelDesc reduced_out_placement = *out->placement();
+  NdSbp reduced_in_nd_sbp;
+  NdSbp reduced_out_nd_sbp;
+  InOutParallelDimReduce(*in->placement(), *out->placement(), *in->nd_sbp(), *out->nd_sbp(),
+                         &reduced_in_placement, &reduced_out_placement, &reduced_in_nd_sbp,
+                         &reduced_out_nd_sbp);
+  return std::make_tuple(
+      JUST(PlacedNdSbp::New(SymbolOf(reduced_in_nd_sbp), SymbolOf(reduced_in_placement))),
+      JUST(PlacedNdSbp::New(SymbolOf(reduced_out_nd_sbp), SymbolOf(reduced_out_placement))));
+}
+
+constexpr auto* InOutPlacedNdSbpDimReduce =
+    DECORATE(&RawInOutPlacedNdSbpDimReduce, ThreadLocalCached);
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckParallelDimReduce(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                      const Shape& logical_shape) {
+  CHECK_OR_RETURN(in->nd_sbp()->sbp_parallel_size() > 1 || out->nd_sbp()->sbp_parallel_size() > 1);
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  Symbol<PlacedNdSbp> reduced_in;
+  Symbol<PlacedNdSbp> reduced_out;
+  std::tie(reduced_in, reduced_out) = *JUST(InOutPlacedNdSbpDimReduce(in, out));
+
+  for (int64_t in_parallel_id = 0; in_parallel_id < in->placement()->parallel_num();
+       ++in_parallel_id) {
+    const auto& in_physical_shape =
+        JUST(GetPhysicalShape(logical_shape, *in->nd_sbp(), *in->placement(), in_parallel_id));
+    const auto& reduce_in_physical_shape = JUST(GetPhysicalShape(
+        logical_shape, *reduced_in->nd_sbp(), *reduced_in->placement(), in_parallel_id));
+    CHECK_EQ_OR_RETURN(*in_physical_shape, *reduce_in_physical_shape);
+  }
+
+  for (int64_t out_parallel_id = 0; out_parallel_id < out->placement()->parallel_num();
+       ++out_parallel_id) {
+    const auto& out_physical_shape =
+        JUST(GetPhysicalShape(logical_shape, *out->nd_sbp(), *out->placement(), out_parallel_id));
+    const auto& reduce_out_physical_shape = JUST(GetPhysicalShape(
+        logical_shape, *reduced_out->nd_sbp(), *reduced_out->placement(), out_parallel_id));
+    CHECK_EQ_OR_RETURN(*out_physical_shape, *reduce_out_physical_shape);
+  }
+
+  if (reduced_in->nd_sbp()->sbp_parallel_size() == 1
+      && reduced_out->nd_sbp()->sbp_parallel_size() == 1) {
+    return Maybe<void>::Ok();
+  }
+  if ((reduced_in->placement() != in->placement() || reduced_out->placement() != out->placement())
+      && reduced_in->placement() == reduced_out->placement()) {
+    return Maybe<void>::Ok();
+  }
+  return Error::CheckFailedError();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckParallelDimReduce =
+    DECORATE(&RawCheckParallelDimReduce, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> ParallelDimReduce(const std::shared_ptr<one::Tensor>& tensor,
+                                     Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+
+  Symbol<PlacedNdSbp> reduced_in;
+  Symbol<PlacedNdSbp> reduced_out;
+  std::tie(reduced_in, reduced_out) = *JUST(InOutPlacedNdSbpDimReduce(in, out));
+
+  const std::shared_ptr<one::Tensor>& local_tensor = JUST(tensor->cur_rank_phy_tensor());
+
+  std::shared_ptr<one::Tensor> reduced_in_tensor = JUST(one::functional::LocalToConsistent(
+      local_tensor, reduced_in->placement(), *JUST(GetSbpList(reduced_in->nd_sbp())),
+      *tensor->shape(), tensor->dtype()));
+
+  const auto& boxing_interpreter =
+      JUST(Singleton<EagerBoxingInterpreterManager>::Get()->GetEagerBoxingInterpreter(
+          reduced_in->nd_sbp(), reduced_out->nd_sbp(), reduced_in->placement(),
+          reduced_out->placement(), *tensor->shape()));
+  Singleton<const EagerBoxingLogger>::Get()->Log(
+      *JUST(boxing_interpreter->boxing_interpreter_status()),
+      /* prefix */ "\t\tInternal boxing of nd-sbp-dim-reduce, ");
+  std::shared_ptr<one::Tensor> reduced_out_tensor = JUST(
+      boxing_interpreter->Interpret(reduced_in_tensor, reduced_in->nd_sbp(), reduced_out->nd_sbp(),
+                                    reduced_in->placement(), reduced_out->placement()));
+
+  const std::shared_ptr<one::Tensor>& reduced_out_local_tensor =
+      JUST(reduced_out_tensor->cur_rank_phy_tensor());
+
+  return JUST(one::functional::LocalToConsistent(reduced_out_local_tensor, out->placement(),
+                                                 *JUST(GetSbpList(out->nd_sbp())), *tensor->shape(),
+                                                 tensor->dtype()));
+}
+
+COMMAND(RegisterBoxingFunction("nd-sbp-dim-reduce", CheckParallelDimReduce, &ParallelDimReduce));
+
+}  // namespace oneflow

oneflow/core/boxing/one_to_one_boxing.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/control/global_process_ctx.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/framework/device.h"
+#include "oneflow/core/boxing/eager_boxing_interpreter.h"
+#include "oneflow/core/functional/functional.h"
+#include "oneflow/core/common/decorator.h"
+
+namespace oneflow {
+
+namespace {
+
+// NOLINTBEGIN(maybe-need-error-msg)
+Maybe<void> RawCheckNaiveOneToOne(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out,
+                                  const Shape& logical_shape) {
+  CHECK_EQ_OR_RETURN(in->placement()->parallel_num(), 1);
+  CHECK_EQ_OR_RETURN(out->placement()->parallel_num(), 1);
+  CHECK_EQ_OR_RETURN(in->placement()->device_tag(), out->placement()->device_tag());
+  CHECK_OR_RETURN(in->placement() != out->placement());
+  return Maybe<void>::Ok();
+}
+// NOLINTEND(maybe-need-error-msg)
+
+static constexpr auto* CheckNaiveOneToOne =
+    DECORATE(&RawCheckNaiveOneToOne, ThreadLocalCachedCopiable);
+
+}  // namespace
+
+Maybe<one::Tensor> NaiveOneToOne(const std::shared_ptr<one::Tensor>& tensor, Symbol<PlacedNdSbp> in,
+                                 Symbol<PlacedNdSbp> out) {
+  const auto& tensor_nd_sbp = JUST(tensor->nd_sbp());
+  CHECK_OR_RETURN(tensor_nd_sbp == in->nd_sbp())
+      << Error::RuntimeError() << "The sbp of input tensor (" << NdSbpToString(tensor_nd_sbp)
+      << ") must match the input sbp (" << NdSbpToString(in->nd_sbp()) << ")";
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == in->placement())
+      << Error::RuntimeError() << "The placement of input tensor ("
+      << *JUST(PlacementToString(tensor_placement)) << ") must match the input placement ("
+      << *JUST(PlacementToString(in->placement())) << ")";
+
+  std::shared_ptr<one::Tensor> local_tensor = JUST(tensor->cur_rank_phy_tensor());
+  int64_t src = JUST(tensor_placement->MachineId4ParallelId(0));
+  int64_t dst = JUST(out->placement()->MachineId4ParallelId(0));
+
+  if (src != dst) {
+    if (GlobalProcessCtx::Rank() == src) {
+      JUST(one::functional::Send(local_tensor, dst, /* send_meta */ false));
+    }
+    if (GlobalProcessCtx::Rank() == dst) {
+      local_tensor = JUST(one::functional::Recv(src, *tensor->shape(), tensor->dtype(),
+                                                JUST(local_tensor->device()), NullOpt));
+    }
+  }
+  return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(),
+                                                 *JUST(GetSbpList(out->nd_sbp())), *tensor->shape(),
+                                                 tensor->dtype()));
+}
+
+COMMAND(RegisterBoxingFunction("naive-1-to-1", CheckNaiveOneToOne, &NaiveOneToOne));
+
+}  // namespace oneflow