Boxing expr

oneflow/core/framework/op_interpreter/boxing/boxing_dividor.h

@@ -0,0 +1,51 @@
+/*
+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_FRAMEWORK_OP_INTERPRETER_BOXING_BOXING_DIVIDOR_H_
+#define ONEFLOW_CORE_FRAMEWORK_OP_INTERPRETER_BOXING_BOXING_DIVIDOR_H_
+
+#include "oneflow/core/common/maybe.h"
+#include "oneflow/core/common/symbol.h"
+
+namespace oneflow {
+
+class PlacedNdSbp;
+
+class BoxingDividor final {
+ public:
+  BoxingDividor(const BoxingDividor&) = delete;
+  BoxingDividor(BoxingDividor&&) = delete;
+  ~BoxingDividor() = default;
+
+  using FunctionT =
+      std::function<Maybe<Symbol<PlacedNdSbp>>(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out)>;
+
+  BoxingDividor(const std::string& name, const FunctionT& function)
+      : name_(name), function_(function) {}
+
+  const std::string& name() const { return name_; }
+
+  Maybe<Symbol<PlacedNdSbp>> operator()(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) const {
+    return function_(in, out);
+  }
+
+ private:
+  std::string name_;
+  FunctionT function_;
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_FRAMEWORK_OP_INTERPRETER_BOXING_BOXING_DIVIDOR_H_

oneflow/core/framework/op_interpreter/boxing/boxing_dividor_util.cpp

@@ -0,0 +1,79 @@
+/*
+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/op_interpreter/boxing/boxing_dividor_util.h"
+#include "oneflow/core/framework/nd_sbp.h"
+#include "oneflow/core/framework/placed_nd_sbp.h"
+#include "oneflow/core/common/decorator.h"
+#include "oneflow/core/job/parallel_desc.h"
+
+namespace oneflow {
+
+namespace {
+
+Maybe<BoxingDividor> RawReplaceInDeviceType(DeviceType device_type) {
+  return std::make_shared<BoxingDividor>(
+      "ReplaceInDeviceType",
+      [device_type](Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) -> Maybe<Symbol<PlacedNdSbp>> {
+        const auto& new_placement = JUST(ReplaceDeviceType(in->placement(), device_type));
+        return PlacedNdSbp::New(in->nd_sbp(), new_placement);
+      });
+}
+
+Maybe<BoxingDividor> RawReplaceOutDeviceType(DeviceType device_type) {
+  return std::make_shared<BoxingDividor>(
+      "ReplaceOutDeviceType",
+      [device_type](Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) -> Maybe<Symbol<PlacedNdSbp>> {
+        const auto& new_placement = JUST(ReplaceDeviceType(out->placement(), device_type));
+        return PlacedNdSbp::New(out->nd_sbp(), new_placement);
+      });
+}
+
+}  // namespace
+
+decltype(ReplaceInDeviceType) ReplaceInDeviceType = DECORATE(&RawReplaceInDeviceType, ThreadLocal);
+decltype(ReplaceOutDeviceType) ReplaceOutDeviceType =
+    DECORATE(&RawReplaceOutDeviceType, ThreadLocal);
+
+namespace {
+
+Maybe<Symbol<PlacedNdSbp>> RawFlattenHierarchy(Symbol<PlacedNdSbp> placed_nd_sbp) {
+  CHECK_GE_OR_RETURN(placed_nd_sbp->nd_sbp()->sbp_parallel_size(), 0);
+  const auto& first_sbp_parallel = placed_nd_sbp->nd_sbp()->sbp_parallel(0);
+  for (const auto& sbp_parallel : placed_nd_sbp->nd_sbp()->sbp_parallel()) {
+    CHECK_OR_RETURN(sbp_parallel == first_sbp_parallel);
+  }
+  std::vector<Symbol<cfg::SbpParallel>> vec{SymbolOf(first_sbp_parallel)};
+  const auto& flattened_nd_sbp = JUST(GetNdSbp(vec));
+  ParallelConf flattened_parallel_conf(placed_nd_sbp->placement()->parallel_conf());
+  flattened_parallel_conf.clear_hierarchy();
+  const auto& flattened_placement = SymbolOf(ParallelDesc(flattened_parallel_conf));
+  return JUST(PlacedNdSbp::New(flattened_nd_sbp, flattened_placement));
+}
+
+static constexpr auto* FlattenHierarchy = DECORATE(&RawFlattenHierarchy, ThreadLocal);
+
+Maybe<BoxingDividor> RawFlattenInHierarchy() {
+  return std::make_shared<BoxingDividor>(
+      "FlattenInHierarchy",
+      [](Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) -> Maybe<Symbol<PlacedNdSbp>> {
+        return FlattenHierarchy(in);
+      });
+}
+
+}  // namespace
+
+decltype(FlattenInHierarchy) FlattenInHierarchy = DECORATE(&RawFlattenInHierarchy, ThreadLocal);
+}  // namespace oneflow

oneflow/core/framework/op_interpreter/boxing/boxing_dividor_util.h

@@ -0,0 +1,30 @@
+/*
+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_FRAMEWORK_OP_INTERPRETER_BOXING_BOXING_DIVIDOR_UTIL_H_
+#define ONEFLOW_CORE_FRAMEWORK_OP_INTERPRETER_BOXING_BOXING_DIVIDOR_UTIL_H_
+
+#include "oneflow/core/common/device_type.pb.h"
+#include "oneflow/core/framework/op_interpreter/boxing/boxing_dividor.h"
+
+namespace oneflow {
+
+extern Maybe<BoxingDividor> (*ReplaceInDeviceType)(DeviceType device_type);
+extern Maybe<BoxingDividor> (*ReplaceOutDeviceType)(DeviceType device_type);
+extern Maybe<BoxingDividor> (*FlattenInHierarchy)();
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_FRAMEWORK_OP_INTERPRETER_BOXING_BOXING_DIVIDOR_UTIL_H_

oneflow/core/framework/op_interpreter/boxing/cuda_copy_boxing_interpreter.cpp

@@ -47,4 +47,45 @@ Maybe<one::Tensor> CudaCopyBoxingInterpreter::InterpretImpl(
   return tensor;
 }
 
+namespace {
+
+Maybe<bool> IgnoringDeviceTypeEqual(Symbol<ParallelDesc> lhs, Symbol<ParallelDesc> rhs) {
+  return lhs == JUST(ReplaceDeviceType(rhs, lhs->device_type()));
+}
+
+}  // namespace
+
+Maybe<void> CheckCudaCopyH2D(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) {
+  bool equal = JUST(IgnoringDeviceTypeEqual(in->placement(), out->placement()));
+  CHECK_OR_RETURN(equal);
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), DeviceType::kCPU);
+  CHECK_EQ_OR_RETURN(out->placement()->device_type(), DeviceType::kGPU);
+  CHECK_OR_RETURN(in->nd_sbp() == out->nd_sbp());
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> CheckCudaCopyD2H(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) {
+  bool equal = JUST(IgnoringDeviceTypeEqual(in->placement(), out->placement()));
+  CHECK_OR_RETURN(equal);
+  CHECK_EQ_OR_RETURN(in->placement()->device_type(), DeviceType::kGPU);
+  CHECK_EQ_OR_RETURN(out->placement()->device_type(), DeviceType::kCPU);
+  CHECK_OR_RETURN(in->nd_sbp() == out->nd_sbp());
+  return Maybe<void>::Ok();
+}
+
+Maybe<one::Tensor> CudaCopy(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());
+  const auto& tensor_placement = JUST(tensor->parallel_desc());
+  CHECK_OR_RETURN(tensor_placement == 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("cuda-copy-h2d", &CheckCudaCopyH2D, &CudaCopy));
+COMMAND(RegisterBoxingFunction("cuda-copy-d2h", &CheckCudaCopyD2H, &CudaCopy));
+
 }  // namespace oneflow

oneflow/core/framework/op_interpreter/boxing/eager_boxing_interpreter.cpp

@@ -83,21 +83,23 @@ HashMap<std::string, BoxingFunctionT>* MutName2BoxingFunction() {
   return &map;
 }
 
-}  // namespace
-
-namespace {
-
 Maybe<BoxingFunctionT> RawGetBoxingFunction(const std::string& method_name, Symbol<PlacedNdSbp> in,
                                             Symbol<PlacedNdSbp> out) {
-  const auto& Checker = JUST(MapAt(*MutName2BoxingChecker(), method_name));
+  const auto& Checker =
+      JUST_MSG(MapAt(*MutName2BoxingChecker(), method_name),
+               std::stringstream() << "boxing checker not found. checker_name: " << method_name);
   JUST(Checker(in, out));
-  return JUST(MapAt(*MutName2BoxingFunction(), method_name));
+  return JUST_MSG(MapAt(*MutName2BoxingFunction(), method_name),
+                  std::stringstream()
+                      << "boxing function not found. function_name: " << method_name);
 }
 
 }  // namespace
 
-decltype(GetBoxingFunction) GetBoxingFunction =
-    DECORATE(&RawGetBoxingFunction, ThreadLocalCopiable);
+Maybe<BoxingFunctionT> GetBoxingFunction(const std::string& method_name, Symbol<PlacedNdSbp> in,
+                                         Symbol<PlacedNdSbp> out) {
+  return DECORATE(&RawGetBoxingFunction, ThreadLocalCopiable)(method_name, in, out);
+}
 
 void RegisterBoxingFunction(const std::string& method_name, const BoxingCheckerT& Checker,
                             const BoxingFunctionT& BoxingFunction) {
@@ -110,4 +112,95 @@ void RegisterBoxingFunction(const std::string& method_name, const BoxingCheckerT
   });
 }
 
+Maybe<void> AtomicBoxingExpr::Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) const {
+  const auto& Checker =
+      JUST_MSG(MapAt(*MutName2BoxingChecker(), boxing_name_),
+               std::stringstream() << "boxing checker not found. checker_name: " << boxing_name_);
+  JUST(Checker(in, out));
+  return Maybe<void>::Ok();
+}
+
+Maybe<BoxingFunctionT> AtomicBoxingExpr::GetBoxingFunction(Symbol<PlacedNdSbp> in,
+                                                           Symbol<PlacedNdSbp> out) const {
+  return DECORATE(&RawGetBoxingFunction, ThreadLocalCopiable)(boxing_name_, in, out);
+}
+
+Maybe<void> DivideAndConquerBoxingExpr::Check(Symbol<PlacedNdSbp> in,
+                                              Symbol<PlacedNdSbp> out) const {
+  const auto& middle = JUST((*boxing_dividor_)(in, out));
+  JUST(lhs_conquer_->Check(in, middle));
+  JUST(rhs_conquer_->Check(middle, out));
+  return Maybe<void>::Ok();
+}
+
+Maybe<BoxingFunctionT> DivideAndConquerBoxingExpr::GetBoxingFunction(
+    Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) const {
+  const auto& middle = JUST((*boxing_dividor_)(in, out));
+  const auto& lhs_boxing_func = JUST(lhs_conquer_->GetBoxingFunction(in, middle));
+  const auto& rhs_boxing_func = JUST(rhs_conquer_->GetBoxingFunction(middle, out));
+  BoxingFunctionT boxing_function = [lhs_boxing_func, rhs_boxing_func, middle, in, out](
+                                        const std::shared_ptr<one::Tensor>& tensor,
+                                        Symbol<PlacedNdSbp> arg_in,
+                                        Symbol<PlacedNdSbp> arg_out) -> Maybe<one::Tensor> {
+    CHECK_OR_RETURN(in == arg_in);
+    CHECK_OR_RETURN(out == arg_out);
+    const auto& middle_tensor = JUST((*lhs_boxing_func)(tensor, in, middle));
+    return JUST((*rhs_boxing_func)(middle_tensor, middle, out));
+  };
+  return boxing_function;
+}
+
+Maybe<void> OrBoxingExpr::Check(Symbol<PlacedNdSbp> in, Symbol<PlacedNdSbp> out) const {
+  if (lhs_boxing_->Check(in, out).IsOk()) { return Maybe<void>::Ok(); }
+  return rhs_boxing_->Check(in, out);
+}
+
+Maybe<BoxingFunctionT> OrBoxingExpr::GetBoxingFunction(Symbol<PlacedNdSbp> in,
+                                                       Symbol<PlacedNdSbp> out) const {
+  if (lhs_boxing_->Check(in, out).IsOk()) { return lhs_boxing_->GetBoxingFunction(in, out); }
+  JUST(rhs_boxing_->Check(in, out));
+  return rhs_boxing_->GetBoxingFunction(in, out);
+}
+
+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