[one::OpBuilder] op interpreter core

oneflow/api/python/framework/foreign_boxing_util.cpp

@@ -52,7 +52,10 @@ class PyForeignBoxingUtil : public ForeignBoxingUtil {
 Maybe<void> RegisterBoxingUtilOnlyOnce(const std::shared_ptr<ForeignBoxingUtil>& boxing_util) {
   CHECK_ISNULL_OR_RETURN(Global<std::shared_ptr<ForeignBoxingUtil>>::Get())
       << "Foreign Boxing util has been registered.";
-  Global<std::shared_ptr<ForeignBoxingUtil>>::New(boxing_util);
+  // Global<T>::SetAllocated is preferred since Global<T>::New will output logs but
+  // glog is not constructed yet.
+  Global<std::shared_ptr<ForeignBoxingUtil>>::SetAllocated(
+      new std::shared_ptr<ForeignBoxingUtil>(boxing_util));
   return Maybe<void>::Ok();
 }
 

oneflow/api/python/framework/op_expr.cpp

@@ -18,33 +18,52 @@ limitations under the License.
 #include "oneflow/api/python/of_api_registry.h"
 #include "oneflow/core/common/protobuf.h"
 #include "oneflow/core/framework/op_expr.h"
+#include "oneflow/core/framework/op_interpreter.h"
+#include "oneflow/core/framework/op_interpreter_util.h"
 #include "oneflow/core/framework/tensor.h"
+#include "oneflow/core/framework/tensor_tuple.h"
 
 namespace py = pybind11;
 
 namespace oneflow {
 
+namespace {
+
+Maybe<one::TensorTuple> Interpret(const one::OpExpr& op, const one::TensorTuple& inputs) {
+  CHECK_EQ_OR_RETURN(op.input_num(), inputs.size())
+      << "The operation requires " << op.input_num() << " inputs, but " << inputs.size()
+      << " is given.";
+  auto outputs = std::make_shared<one::TensorTuple>(op.output_num());
+  auto interperter = JUST(one::OpInterpUtil::GetInterpreter());
+  JUST(interperter->Apply(op, inputs, outputs.get()));
+  return outputs;
+}
+
 Maybe<std::vector<std::shared_ptr<one::Tensor>>> Interpret(
-    const std::shared_ptr<one::OpExpr>& op,
-    const std::vector<std::shared_ptr<one::Tensor>>& inputs) {
-  // TODO(): Execute the op by Autograd.
-  UNIMPLEMENTED();
-  return std::make_shared<std::vector<std::shared_ptr<one::Tensor>>>();
+    const one::OpExpr& op, const std::vector<std::shared_ptr<one::Tensor>>& inputs) {
+  one::TensorTuple input_list(inputs.size());
+  for (int i = 0; i < inputs.size(); ++i) { input_list[i] = inputs[i]; }
+  const auto& outputs = JUST(Interpret(op, input_list));
+  return static_cast<std::shared_ptr<std::vector<std::shared_ptr<one::Tensor>>>>(outputs);
 }
 
+}  // namespace
+
 ONEFLOW_API_PYBIND11_MODULE("one", m) {
   py::class_<one::OpExpr, std::shared_ptr<one::OpExpr>>(m, "OpExpr")
-      .def("apply", [](const std::shared_ptr<one::OpExpr>& op_expr,
-                       const std::vector<std::shared_ptr<one::Tensor>>& inputs) {
-        return Interpret(op_expr, inputs).GetOrThrow();
+      .def("apply",
+           [](const one::OpExpr& op_expr, const std::vector<std::shared_ptr<one::Tensor>>& inputs) {
+             return Interpret(op_expr, inputs).GetOrThrow();
+           })
+      .def("apply", [](const one::OpExpr& op_expr, const one::TensorTuple& inputs) {
+        return Interpret(op_expr, inputs).GetPtrOrThrow();
       });
 
   py::class_<one::BuiltinOpExpr, one::OpExpr, std::shared_ptr<one::BuiltinOpExpr>>(m,
                                                                                    "BuiltinOpExpr")
       .def_property_readonly("name", &one::BuiltinOpExpr::op_name);
 
   py::class_<one::UserOpExpr, one::BuiltinOpExpr, std::shared_ptr<one::UserOpExpr>>(m, "UserOpExpr")
-      .def(py::init<>())
       .def_property_readonly("type", &one::UserOpExpr::type)
       .def_property_readonly(
           "proto", [](const one::UserOpExpr& op) { return PbMessage2TxtString(op.proto()); })

oneflow/api/python/framework/session_util.cpp

@@ -23,7 +23,19 @@ namespace oneflow {
 
 ONEFLOW_API_PYBIND11_MODULE("", m) {
   py::class_<Session, std::shared_ptr<Session>>(m, "Session")
-      .def_property_readonly("id", &Session::id);
+      .def_property_readonly("id", &Session::id)
+      .def("push_mirrored_strategy_enabled",
+           [](Session* sess, const bool& val) {
+             sess->PushMirroredStrategyEnabled(val).GetOrThrow();
+           })
+      .def("pop_mirrored_strategy_enabled",
+           [](Session* sess) { sess->PopMirroredStrategyEnabled().GetOrThrow(); })
+      .def("is_mirrored_strategy_enabled",
+           [](const Session* sess) { return sess->IsMirroredStrategyEnabled().GetOrThrow(); })
+      .def("is_consistent_strategy_enabled",
+           [](const Session* sess) { return sess->IsConsistentStrategyEnabled().GetOrThrow(); })
+      .def("is_mirrored_strategy_enabled_stack_size",
+           [](const Session* sess) { return sess->is_mirrored_strategy_enabled_stack()->size(); });
 
   m.def("GetDefaultSessionId", []() { return GetDefaultSessionId().GetOrThrow(); });
   m.def("RegsiterSession", [](int64_t id) { return RegsiterSession(id).GetPtrOrThrow(); });

oneflow/core/framework/op_builder.cpp

@@ -38,10 +38,10 @@ Maybe<OpBuilder&> OpBuilder::MaybeInput(const std::string& input_name, const int
   CHECK_GT_OR_RETURN(count, 0);
   CHECK_EQ_OR_RETURN(proto_.input().count(input_name), 0)
       << "The Input " << input_name << " has been specified more than once.";
-  auto& input_list = (*(proto_.mutable_input()))[input_name];
+  auto* input_list = &((*(proto_.mutable_input()))[input_name]);
   for (int i = 0; i < count; ++i) {
     const std::string& tensor_name = _PositionalPlaceholderPrefix + std::to_string(input_pos_++);
-    input_list.mutable_s()->Add()->assign(tensor_name);
+    input_list->mutable_s()->Add()->assign(tensor_name);
     indexed_ibns_.push_back(input_name + "_" + std::to_string(i));
   }
   return *this;
@@ -58,10 +58,10 @@ Maybe<OpBuilder&> OpBuilder::MaybeOutput(const std::string& output_name, const i
   CHECK_GT_OR_RETURN(count, 0);
   CHECK_EQ_OR_RETURN(proto_.output().count(output_name), 0)
       << "The output " << output_name << " has been specified more than once.";
-  auto& output_list = (*(proto_.mutable_output()))[output_name];
+  auto* output_list = &((*(proto_.mutable_output()))[output_name]);
   for (int i = 0; i < count; ++i) {
     const std::string& tensor_name = op_name_ + "/" + output_name + "_" + std::to_string(i);
-    output_list.mutable_s()->Add()->assign(tensor_name);
+    output_list->mutable_s()->Add()->assign(tensor_name);
     indexed_obns_.push_back(output_name + "_" + std::to_string(i));
   }
   return *this;

oneflow/core/framework/op_expr.h

@@ -16,37 +16,40 @@ limitations under the License.
 #ifndef ONEFLOW_CORE_FRAMEWORK_OP_EXPR_H_
 #define ONEFLOW_CORE_FRAMEWORK_OP_EXPR_H_
 
+#include "oneflow/core/common/util.h"
 #include "oneflow/core/framework/user_op_conf.pb.h"
 #include "oneflow/core/operator/op_conf.pb.h"
 
 namespace oneflow {
 namespace one {
 
-#define DEFINE_DEFAULT_CONSTRUCTOR(class_type) \
-  class_type() = default;                      \
-  virtual ~class_type() = default;
-
 class OpExpr {
  public:
-  DEFINE_DEFAULT_CONSTRUCTOR(OpExpr);
+  explicit OpExpr(const std::string& type) : type_(type) {}
+  virtual ~OpExpr() = default;
+
+  const std::string& type() const { return type_; }
 
-  virtual std::shared_ptr<OpExpr> GetBackwardOpExpr() const = 0;
-  virtual std::string type() const = 0;
+  virtual int input_num() const = 0;
+  virtual int output_num() const = 0;
+
+ private:
+  std::string type_;
 };
 
 class BuiltinOpExpr : public OpExpr {
  public:
-  DEFINE_DEFAULT_CONSTRUCTOR(BuiltinOpExpr);
+  explicit BuiltinOpExpr(const std::string& type, const std::string& op_name,
+                         const std::vector<std::string>& indexed_ibns,
+                         const std::vector<std::string>& indexed_obns)
+      : OpExpr(type), op_name_(op_name), indexed_ibns_(indexed_ibns), indexed_obns_(indexed_obns) {}
 
-  explicit BuiltinOpExpr(const std::string& op_name) : op_name_(op_name) {}
-  BuiltinOpExpr(const std::string& op_name, const std::vector<std::string>& indexed_ibns,
-                const std::vector<std::string>& indexed_obns)
-      : op_name_(op_name), indexed_ibns_(indexed_ibns), indexed_obns_(indexed_obns) {}
+  virtual ~BuiltinOpExpr() = default;
 
   const std::string& op_name() const { return op_name_; }
 
-  int input_num() const { return indexed_ibns_.size(); }
-  int output_num() const { return indexed_obns_.size(); }
+  int input_num() const override { return indexed_ibns_.size(); }
+  int output_num() const override { return indexed_obns_.size(); }
 
   const std::vector<std::string>& indexed_ibns() const { return indexed_ibns_; }
   const std::vector<std::string>& indexed_obns() const { return indexed_obns_; }
@@ -61,30 +64,27 @@ class BuiltinOpExpr : public OpExpr {
   std::vector<std::string> indexed_obns_;
 };
 
-#define DEFINE_BUILTIN_OPEXPR_CLASS(_op_name, _op_conf)                         \
-  class _op_name##Expr : public BuiltinOpExpr {                                 \
-   public:                                                                      \
-    _op_name##Expr() = default;                                                 \
-    virtual ~_op_name##Expr() = default;                                        \
-    explicit _op_name##Expr(const std::string& op_name, _op_name##Conf&& proto, \
-                            const std::vector<std::string>& indexed_ibns,       \
-                            const std::vector<std::string>& indexed_obns)       \
-        : BuiltinOpExpr(op_name, indexed_ibns, indexed_obns), proto_(proto) {}  \
-                                                                                \
-    std::shared_ptr<OpExpr> GetBackwardOpExpr() const override;                 \
-                                                                                \
-    std::string type() const override { return std::string(#_op_name); }        \
-                                                                                \
-    const _op_name##Conf& proto() const { return proto_; }                      \
-    _op_name##Conf* mutable_proto() { return &proto_; }                         \
-                                                                                \
-    void BuildOpConf(OperatorConf* op_conf) const {                             \
-      *(op_conf->mutable_name()) = this->op_name_;                              \
-      *(op_conf->mutable_##_op_conf##_conf()) = proto_;                         \
-    }                                                                           \
-                                                                                \
-   private:                                                                     \
-    _op_name##Conf proto_;                                                      \
+#define DEFINE_BUILTIN_OPEXPR_CLASS(_op_name, _op_conf)                                  \
+  class _op_name##Expr : public BuiltinOpExpr {                                          \
+   public:                                                                               \
+    _op_name##Expr() = default;                                                          \
+    virtual ~_op_name##Expr() = default;                                                 \
+    explicit _op_name##Expr(const std::string& op_name, _op_name##Conf&& proto,          \
+                            const std::vector<std::string>& indexed_ibns,                \
+                            const std::vector<std::string>& indexed_obns)                \
+        : BuiltinOpExpr(OF_PP_STRINGIZE(_op_name), op_name, indexed_ibns, indexed_obns), \
+          proto_(std::move(proto)) {}                                                    \
+                                                                                         \
+    const _op_name##Conf& proto() const { return proto_; }                               \
+    _op_name##Conf* mutable_proto() { return &proto_; }                                  \
+                                                                                         \
+    void BuildOpConf(OperatorConf* op_conf) const {                                      \
+      *(op_conf->mutable_name()) = this->op_name_;                                       \
+      *(op_conf->mutable_##_op_conf##_conf()) = proto_;                                  \
+    }                                                                                    \
+                                                                                         \
+   private:                                                                              \
+    _op_name##Conf proto_;                                                               \
   };
 
 DEFINE_BUILTIN_OPEXPR_CLASS(UserOp, user);
@@ -101,16 +101,16 @@ DEFINE_BUILTIN_OPEXPR_CLASS(DistributeAddOp, distribute_add);
 // TODO(): Finish the class definition of `FunctionOpExpr`.
 class FunctionOpExpr : public OpExpr {
  public:
-  DEFINE_DEFAULT_CONSTRUCTOR(FunctionOpExpr);
-
-  std::shared_ptr<OpExpr> GetBackwardOpExpr() const override;
+  FunctionOpExpr() : OpExpr("FunctionOp") {}
+  virtual ~FunctionOpExpr() = default;
 
-  std::string type() const override { return "FunctionOp"; }
+  int input_num() const override { UNIMPLEMENTED(); }
+  int output_num() const override { UNIMPLEMENTED(); }
 };
 
-#undef DEFINE_DEFAULT_CONSTRUCTOR
-
 }  // namespace one
 }  // namespace oneflow
 
+#undef DEFINE_DEFAULT_CONSTRUCTOR
+
 #endif  // ONEFLOW_CORE_FRAMEWORK_OP_EXPR_H_