[Random Generator] Part2: Migrate functional dropout (#5378)

* add random generator

* reformat

* refactor: allow auto generator

* refactor: remove kAUTO, update python api and test

* refactor: use member initializer lists, fix build issue when cpu only

* handle exeception given invalid device

* add dropout functor; add OpExprInterpContext; refactor random_mask_like based on random_generator

* refactor random generator

* disable generator's copyandmove constructor

* reformat

* fix bad merge

* refine

* fix cpu only build

* auto format by CI

* refactor

* use global generator when no generator specified in functional api

* refine

Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org>
Co-authored-by: Houjiang Chen <chenhoujiangcug@gmail.com>
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/api/python/functional/python_arg.cpp

@@ -23,6 +23,7 @@ limitations under the License.
 #include "oneflow/core/framework/tensor.h"
 #include "oneflow/core/framework/tensor_tuple.h"
 #include "oneflow/core/framework/user_op_attr.cfg.h"
+#include "oneflow/core/framework/random_generator.h"
 #include "oneflow/core/functional/scalar.h"
 
 namespace py = pybind11;
@@ -155,6 +156,17 @@ Maybe<Shape> PythonArg::ObjectAs<Shape>() const {
   }
 }
 
+template<>
+Maybe<std::shared_ptr<one::Generator>> PythonArg::ObjectAs<std::shared_ptr<one::Generator>>()
+    const {
+  return detail::cast<std::shared_ptr<one::Generator>>(Borrow());
+}
+
+template<>
+Maybe<one::Generator> PythonArg::ObjectAs<one::Generator>() const {
+  return *JUST(detail::cast<std::shared_ptr<one::Generator>>(Borrow()));
+}
+
 }  // namespace functional
 }  // namespace one
 }  // namespace oneflow

oneflow/core/autograd/gradient_funcs/adaptive_pool.cpp

@@ -64,7 +64,7 @@ Maybe<void> AdaptivePool::Apply(const AdaptivePoolInterpState* ctx, const Tensor
 
   const std::shared_ptr<oneflow::one::Tensor>& x = ctx->SavedTensors().at(0);
   in_grads->resize(1);
-  in_grads->at(0) = JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {x, out_grads.at(0)}, {}));
+  in_grads->at(0) = JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {x, out_grads.at(0)}));
   return Maybe<void>::Ok();
 }
 

oneflow/core/autograd/gradient_funcs/bias_add.cpp

@@ -68,8 +68,8 @@ class BiasAdd : public OpExprGradFunction<BiasAddInterpState> {
           JUST(OpInterpUtil::Dispatch<Tensor>(*backward_bias_op_, {out_grads.at(0)}, attrs));
     }
     if (ctx->input_requires_grad) {
-      in_grads->at(0) = JUST(
-          OpInterpUtil::Dispatch<Tensor>(*backward_input_op_, {out_grads.at(0)}, /*attrs=*/{}));
+      in_grads->at(0) =
+          JUST(OpInterpUtil::Dispatch<Tensor>(*backward_input_op_, {out_grads.at(0)}));
     }
     return Maybe<void>::Ok();
   }

oneflow/core/autograd/gradient_funcs/broadcast_binary_ops.cpp

@@ -134,8 +134,7 @@ class BroadcastSub : public BroadcastBinaryGrad {
     in_grads->resize(2);
     if (x->requires_grad()) { in_grads->at(0) = JUST(x_grad_op_->forward(out_grads.at(0), x)); }
     if (y->requires_grad()) {
-      const auto& grad =
-          JUST(OpInterpUtil::Dispatch<Tensor>(*y_grad_mul_op_, {out_grads.at(0)}, /*attrs=*/{}));
+      const auto& grad = JUST(OpInterpUtil::Dispatch<Tensor>(*y_grad_mul_op_, {out_grads.at(0)}));
       in_grads->at(1) = JUST(y_grad_op_->forward(grad, y));
     }
     return Maybe<void>::Ok();
@@ -169,12 +168,12 @@ class BroadcastMul : public BroadcastBinaryGrad {
     in_grads->resize(2);
     if (x->requires_grad()) {
       const auto& x_grad =
-          JUST(OpInterpUtil::Dispatch<Tensor>(*x_grad_mul_op_, {out_grads.at(0), y}, /*attrs=*/{}));
+          JUST(OpInterpUtil::Dispatch<Tensor>(*x_grad_mul_op_, {out_grads.at(0), y}));
       in_grads->at(0) = JUST(x_grad_op_->forward(x_grad, x));
     }
     if (y->requires_grad()) {
       const auto& y_grad =
-          JUST(OpInterpUtil::Dispatch<Tensor>(*y_grad_mul_op_, {out_grads.at(0), x}, /*attrs=*/{}));
+          JUST(OpInterpUtil::Dispatch<Tensor>(*y_grad_mul_op_, {out_grads.at(0), x}));
       in_grads->at(1) = JUST(y_grad_op_->forward(y_grad, y));
     }
     return Maybe<void>::Ok();
@@ -208,12 +207,11 @@ class BroadcastDiv : public BroadcastBinaryGrad {
     in_grads->resize(2);
     if (x->requires_grad()) {
       const auto& x_grad =
-          JUST(OpInterpUtil::Dispatch<Tensor>(*x_grad_div_op_, {out_grads.at(0), y}, /*attrs=*/{}));
+          JUST(OpInterpUtil::Dispatch<Tensor>(*x_grad_div_op_, {out_grads.at(0), y}));
       in_grads->at(0) = JUST(x_grad_op_->forward(x_grad, x));
     }
     if (y->requires_grad()) {
-      in_grads->at(1) =
-          JUST(OpInterpUtil::Dispatch<Tensor>(*y_grad_op_, {out_grads.at(0), z, y}, /*attrs=*/{}));
+      in_grads->at(1) = JUST(OpInterpUtil::Dispatch<Tensor>(*y_grad_op_, {out_grads.at(0), z, y}));
     }
     return Maybe<void>::Ok();
   }

oneflow/core/autograd/gradient_funcs/deconv.cpp

@@ -98,8 +98,7 @@ Maybe<void> DeConvolutionNd::Apply(const DeConvolutionNdInterpState* ctx,
   if (ctx->weight_requires_grad) {
     int idx = ctx->activation_requires_grad;
     const auto& x = ctx->SavedTensors().at(idx);
-    in_grads->at(1) =
-        JUST(OpInterpUtil::Dispatch<Tensor>(*weight_grad_op_, {x, out_grads.at(0)}, /*attrs=*/{}));
+    in_grads->at(1) = JUST(OpInterpUtil::Dispatch<Tensor>(*weight_grad_op_, {x, out_grads.at(0)}));
   }
   return Maybe<void>::Ok();
 }

oneflow/core/autograd/gradient_funcs/squeeze.cpp

@@ -64,8 +64,7 @@ Maybe<void> Squeeze::Apply(const SqueezeInterpState* ctx, const TensorTuple& out
 
   const std::shared_ptr<oneflow::one::Tensor>& like = ctx->SavedTensors().at(0);
   in_grads->resize(1);
-  in_grads->at(0) =
-      JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {out_grads.at(0), like}, /*attrs*/ {}));
+  in_grads->at(0) = JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {out_grads.at(0), like}));
   return Maybe<void>::Ok();
 }
 

oneflow/core/autograd/gradient_funcs/tensor_scalar_binary.cpp

@@ -89,8 +89,7 @@ class TensorScalarSub : public TensorScalarAddOrSub {
                     TensorTuple* in_grads) const override {
     in_grads->resize(2);
     if (ctx->x_requires_grad) {
-      in_grads->at(0) =
-          JUST(OpInterpUtil::Dispatch<Tensor>(*identity_op_, {out_grads.at(0)}, /*attrs=*/{}));
+      in_grads->at(0) = JUST(OpInterpUtil::Dispatch<Tensor>(*identity_op_, {out_grads.at(0)}));
     }
     if (ctx->scalar_requires_grad) {
       int32_t num_axes = out_grads.at(0)->shape()->NumAxes();

oneflow/core/autograd/gradient_funcs/unsqueeze.cpp

@@ -64,8 +64,7 @@ Maybe<void> Unsqueeze::Apply(const UnsqueezeInterpState* ctx, const TensorTuple&
 
   const std::shared_ptr<oneflow::one::Tensor>& like = ctx->SavedTensors().at(0);
   in_grads->resize(1);
-  in_grads->at(0) =
-      JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {out_grads.at(0), like}, /*attrs*/ {}));
+  in_grads->at(0) = JUST(OpInterpUtil::Dispatch<Tensor>(*grad_op_, {out_grads.at(0), like}));
   return Maybe<void>::Ok();
 }
 

oneflow/core/eager/eager_oneflow.cpp

@@ -68,11 +68,11 @@ Maybe<void> StorageAdd(const EagerSymbol& symbol) {
 Maybe<void> EagerOneflow::RunPhysicalInstruction(
     const std::shared_ptr<const ClusterInstructionProto>& cluster_instruction) {
   vm::InstructionMsgList instruction_list;
-  const auto& eage_instructions = cluster_instruction->eager_instruction();
-  for (const auto& instr_proto : eage_instructions.instruction_list().instruction()) {
+  const auto& eager_instructions = cluster_instruction->eager_instruction();
+  for (const auto& instr_proto : eager_instructions.instruction_list().instruction()) {
     instruction_list.EmplaceBack(ObjectMsgPtr<vm::InstructionMsg>::New(instr_proto));
   }
-  return RunPhysicalInstruction(&instruction_list, eage_instructions.eager_symbol_list());
+  return RunPhysicalInstruction(&instruction_list, eager_instructions.eager_symbol_list());
 }
 
 Maybe<void> EagerOneflow::RunPhysicalInstruction(

oneflow/core/eager/local_call_opkernel_phy_instr_operand.h

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "oneflow/core/eager/eager_blob_object.h"
 #include "oneflow/core/framework/attr_map.h"
+#include "oneflow/core/framework/op_interpreter.h"
 #include "oneflow/core/vm/instruction_operand.msg.h"
 
 namespace oneflow {
@@ -48,13 +49,15 @@ class LocalCallOpKernelPhyInstrOperand final : public vm::PhyInstrOperand {
 
   LocalCallOpKernelPhyInstrOperand(const std::shared_ptr<one::StatefulLocalOpKernel>& opkernel,
                                    const one::EagerBlobObjectListPtr& inputs,
-                                   const one::EagerBlobObjectListPtr& outputs, const AttrMap& attrs)
-      : opkernel_(opkernel), inputs_(inputs), outputs_(outputs), attrs_(attrs) {}
+                                   const one::EagerBlobObjectListPtr& outputs,
+                                   const one::OpExprInterpContext& op_interp_ctx_)
+      : opkernel_(opkernel), inputs_(inputs), outputs_(outputs), op_interp_ctx_(op_interp_ctx_) {}
 
   const one::StatefulLocalOpKernel& opkernel() const { return *opkernel_; }
   const one::EagerBlobObjectListPtr& inputs() const { return inputs_; }
   const one::EagerBlobObjectListPtr& outputs() const { return outputs_; }
-  const AttrMap& attrs() const { return attrs_; }
+  const AttrMap& attrs() const { return op_interp_ctx_.attrs; }
+  const one::OpExprInterpContext& op_interp_ctx() const { return op_interp_ctx_; }
 
   one::StatefulLocalOpKernel* mut_opkernel() { return opkernel_.get(); }
 
@@ -84,7 +87,7 @@ class LocalCallOpKernelPhyInstrOperand final : public vm::PhyInstrOperand {
   std::shared_ptr<one::StatefulLocalOpKernel> opkernel_;
   one::EagerBlobObjectListPtr inputs_;
   one::EagerBlobObjectListPtr outputs_;
-  const AttrMap attrs_;
+  const one::OpExprInterpContext op_interp_ctx_;
   const user_op::OpKernel* user_opkernel_;
 };
 

oneflow/core/eager/opkernel_instruction_type.cpp

@@ -550,6 +550,10 @@ struct LocalCallOpKernelUtil final {
 
   static inline void TryInitOpKernelState(LocalCallOpKernelPhyInstrOperand* operand,
                                           DeviceCtx* device_ctx, user_op::OpKernelState** state) {
+    if (operand->op_interp_ctx().state) {
+      *state = operand->op_interp_ctx().state.get();
+      return;
+    }
     operand->mut_opkernel()->TryInitOpKernelState(operand->user_opkernel(), device_ctx,
                                                   operand->inputs(), operand->outputs(), state);
   }

oneflow/core/framework/instructions_builder.cpp

@@ -653,13 +653,14 @@ Maybe<void> InstructionsBuilder::BuildRecvInstruction(
 Maybe<void> InstructionsBuilder::LocalCallOpKernel(
     const std::shared_ptr<one::StatefulLocalOpKernel>& opkernel,
     const one::EagerBlobObjectListPtr& input_eager_blob_objects,
-    const one::EagerBlobObjectListPtr& output_eager_blob_objects, const AttrMap& attrs,
+    const one::EagerBlobObjectListPtr& output_eager_blob_objects,
+    const one::OpExprInterpContext& ctx,
     const std::shared_ptr<const ParallelDesc>& parallel_desc_sym,
     const std::string& instr_type_name) {
   ObjectMsgPtr<vm::InstructionMsg> instruction =
       ObjectMsgPtr<vm::InstructionMsg>::New(instr_type_name);
   auto phy_instr_operand = std::make_shared<vm::LocalCallOpKernelPhyInstrOperand>(
-      opkernel, input_eager_blob_objects, output_eager_blob_objects, attrs);
+      opkernel, input_eager_blob_objects, output_eager_blob_objects, ctx);
   *instruction->mut_parallel_desc() = parallel_desc_sym;
   *instruction->mutable_phy_instr_operand() = phy_instr_operand;
   instruction_list_->EmplaceBack(std::move(instruction));

oneflow/core/framework/instructions_builder.h

@@ -247,7 +247,7 @@ class InstructionsBuilder : public std::enable_shared_from_this<InstructionsBuil
   Maybe<void> LocalCallOpKernel(const std::shared_ptr<one::StatefulLocalOpKernel>& opkernel,
                                 const one::EagerBlobObjectListPtr& input_eager_blob_objects,
                                 const one::EagerBlobObjectListPtr& output_eager_blob_objects,
-                                const AttrMap& attrs,
+                                const one::OpExprInterpContext& ctx,
                                 const std::shared_ptr<const ParallelDesc>& parallel_desc_sym,
                                 const std::string& instr_type_name);
 

oneflow/core/framework/op_interpreter.h

@@ -20,6 +20,7 @@ limitations under the License.
 #include "oneflow/core/framework/op_expr.h"
 #include "oneflow/core/framework/tensor.h"
 #include "oneflow/core/framework/tensor_tuple.h"
+#include "oneflow/core/framework/op_kernel.h"
 
 namespace oneflow {
 namespace one {
@@ -41,17 +42,27 @@ class OpExprInterpState {
   TensorTuple saved_tensors_;
 };
 
+struct OpExprInterpContext {
+  AttrMap attrs;
+  std::shared_ptr<user_op::OpKernelState> state;
+};
+
 class OpExprInterpreter {
  public:
   OpExprInterpreter() = default;
   virtual ~OpExprInterpreter() = default;
 
-  virtual Maybe<void> Apply(const OpExpr& op, const TensorTuple& inputs, TensorTuple* outputs,
-                            const AttrMap& attrs) const = 0;
+  Maybe<void> Apply(const OpExpr& op, const TensorTuple& inputs, TensorTuple* outputs,
+                    const AttrMap& attrs) const {
+    return Apply(op, inputs, outputs, OpExprInterpContext{attrs, nullptr});
+  }
 
   Maybe<void> Apply(const OpExpr& op, const TensorTuple& inputs, TensorTuple* outputs) const {
     return Apply(op, inputs, outputs, AttrMap{});
   }
+
+  virtual Maybe<void> Apply(const OpExpr& op, const TensorTuple& inputs, TensorTuple* outputs,
+                            const OpExprInterpContext& ctx) const = 0;
 };
 
 #define FOR_EACH_BUILTIN_OPS(_macro) \
@@ -66,21 +77,26 @@ class OpExprInterpreter {
 
 #define DECLARE_NORMAL_APPLY_FUNC(op_type)                                               \
   virtual Maybe<void> ApplyImpl(const op_type##Expr& op_expr, const TensorTuple& inputs, \
-                                TensorTuple* outputs, const AttrMap& attrs) const
+                                TensorTuple* outputs, const OpExprInterpContext& ctx) const
 
 #define DECLARE_PURE_VIRTUAL_APPLY_FUNC(op_type) DECLARE_NORMAL_APPLY_FUNC(op_type) = 0;
 
 #define DECLARE_OVERRIDE_APPLY_FUNC(op_type)                                     \
   Maybe<void> ApplyImpl(const op_type##Expr& op_expr, const TensorTuple& inputs, \
-                        TensorTuple* outputs, const AttrMap& attrs) const override;
+                        TensorTuple* outputs, const OpExprInterpContext& ctx) const override;
 
 class LazyInterpreter : public OpExprInterpreter {
  public:
   LazyInterpreter() : OpExprInterpreter() {}
   virtual ~LazyInterpreter() = default;
 
   Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs,
-                    const AttrMap& attrs) const override;
+                    const AttrMap& attrs) const {
+    return Apply(op_expr, inputs, outputs, OpExprInterpContext{attrs, nullptr});
+  }
+
+  Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs,
+                    const OpExprInterpContext& ctx) const override;
 
  private:
   DECLARE_NORMAL_APPLY_FUNC(BuiltinOp);
@@ -93,7 +109,12 @@ class EagerInterpreter : public OpExprInterpreter {
   virtual ~EagerInterpreter() = default;
 
   Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs,
-                    const AttrMap& attrs) const override;
+                    const AttrMap& attrs) const {
+    return Apply(op_expr, inputs, outputs, OpExprInterpContext{attrs, nullptr});
+  }
+
+  Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs,
+                    const OpExprInterpContext& ctx) const override;
 
  private:
   FOR_EACH_BUILTIN_OPS(DECLARE_PURE_VIRTUAL_APPLY_FUNC);
@@ -131,12 +152,17 @@ class AutogradInterpreter {
   virtual ~AutogradInterpreter() = default;
 
   Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs,
-                    const AttrMap& attrs) const;
+                    const AttrMap& attrs) const {
+    return Apply(op_expr, inputs, outputs, OpExprInterpContext{attrs, nullptr});
+  }
 
-  Maybe<void> Apply(const OpExpr& op, const TensorTuple& inputs, TensorTuple* outputs) const {
-    return Apply(op, inputs, outputs, AttrMap{});
+  Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs) const {
+    return Apply(op_expr, inputs, outputs, OpExprInterpContext{});
   }
 
+  Maybe<void> Apply(const OpExpr& op_expr, const TensorTuple& inputs, TensorTuple* outputs,
+                    const OpExprInterpContext& ctx) const;
+
  private:
   std::shared_ptr<OpExprInterpreter> internal_;
 };

oneflow/core/framework/op_interpreter/eager_consistent_op_interpreter.cpp

@@ -33,10 +33,11 @@ namespace oneflow {
 namespace one {
 
 Maybe<void> Interpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,
-                      TensorTuple* outputs, const AttrMap& attrs) {
+                      TensorTuple* outputs, const OpExprInterpContext& ctx) {
   CHECK_EQ_OR_RETURN(outputs->size(), user_op_expr.output_size());
   const auto& placement_scope = JUST(GetCurrentScope())->placement_scope();
-  const auto& infer_args = JUST(ConsistentTensorMetaInferArgs::New(inputs, placement_scope, attrs));
+  const auto& infer_args =
+      JUST(ConsistentTensorMetaInferArgs::New(inputs, placement_scope, ctx.attrs));
   const auto& result =
       JUST(user_op_expr.mut_consistent_tensor_infer_cache()->GetOrInfer(*infer_args));
   const auto& output_tensor_metas = result->output_tensor_metas();
@@ -71,56 +72,56 @@ Maybe<void> Interpret(const UserOpExpr& user_op_expr, const TensorTuple& inputs,
   const auto& instr_type_name = JUST(GetLocalCallInstructionName(parallel_desc->device_tag()));
   JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
     return builder->LocalCallOpKernel(kernel, input_eager_blob_objects, output_eager_blob_objects,
-                                      attrs, parallel_desc, instr_type_name);
+                                      ctx, parallel_desc, instr_type_name);
   }));
   return Maybe<void>::Ok();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const UserOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
-  return Interpret(op_expr, inputs, outputs, attrs);
+                                                  const OpExprInterpContext& ctx) const {
+  OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const VariableOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const CastToMirroredOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const CastFromMirroredOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const DistributeSplitOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const DistributeCloneOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const DistributeConcatOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }
 
 Maybe<void> EagerConsistentInterpreter::ApplyImpl(const DistributeAddOpExpr& op_expr,
                                                   const TensorTuple& inputs, TensorTuple* outputs,
-                                                  const AttrMap& attrs) const {
+                                                  const OpExprInterpContext& ctx) const {
   OF_UNIMPLEMENTED();
 }