add some fused kernels

oneflow/core/autograd/gradient_funcs/fused_scale_mask_softmax.cpp

@@ -0,0 +1,85 @@
+/*
+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_expr_grad_function.h"
+#include "oneflow/core/framework/op_builder.h"
+#include "oneflow/core/framework/op_expr.h"
+#include "oneflow/core/framework/op_expr_helper.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+struct FusedScaleMaskSoftmaxInterState : public AutoGradCaptureState {
+  bool input_requires_grad = false;
+  float scale = 1.0;
+};
+
+class FusedScaleMaskSoftmax : public OpExprGradFunction<FusedScaleMaskSoftmaxInterState> {
+ public:
+  Maybe<void> Init(const OpExpr& op) override;
+  Maybe<void> Capture(FusedScaleMaskSoftmaxInterState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const FusedScaleMaskSoftmaxInterState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override;
+
+ private:
+  AttrMap base_attrs_;
+};
+
+Maybe<void> FusedScaleMaskSoftmax::Init(const OpExpr& op) {
+  const UserOpExpr* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
+  CHECK_NOTNULL_OR_RETURN(fw_op_expr);
+  base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> FusedScaleMaskSoftmax::Capture(FusedScaleMaskSoftmaxInterState* ctx,
+                                           const TensorTuple& inputs, const TensorTuple& outputs,
+                                           const AttrMap& attrs) const {
+  CHECK_EQ_OR_RETURN(inputs.size(), 2); // input, mask
+  ctx->input_requires_grad = inputs.at(0)->requires_grad();
+
+  if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
+  ComposedAttrMap composed_attrs(attrs, base_attrs_);
+  ctx->scale = JUST(composed_attrs.GetAttr<float>("scale_value"));
+
+  ctx->SaveTensorForBackward(inputs.at(1));  // save mask
+  ctx->SaveTensorForBackward(outputs.at(0)); // save y, ie. softmax result
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> FusedScaleMaskSoftmax::Apply(const FusedScaleMaskSoftmaxInterState* ctx,
+                                         const TensorTuple& out_grads, TensorTuple* in_grads) const {
+  if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
+
+  CHECK_EQ_OR_RETURN(out_grads.size(), 1);  // dy
+  in_grads->resize(2);  // input, mask
+
+  const std::shared_ptr<oneflow::one::Tensor>& mask = ctx->SavedTensors().at(0);
+  const std::shared_ptr<oneflow::one::Tensor>& y = ctx->SavedTensors().at(1);
+  const std::shared_ptr<oneflow::one::Tensor>& fused_scale_mask_softmax_grad = 
+      JUST(functional::FusedScaleMaskSoftmaxGrad(y, out_grads.at(0), mask, ctx->scale));
+
+  in_grads->at(0) = fused_scale_mask_softmax_grad;
+  return Maybe<void>::Ok();
+}
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("fused_scale_mask_softmax", FusedScaleMaskSoftmax);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/core/autograd/gradient_funcs/fused_scale_mask_softmax_dropout.cpp

@@ -0,0 +1,89 @@
+/*
+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_expr_grad_function.h"
+#include "oneflow/core/framework/op_builder.h"
+#include "oneflow/core/framework/op_expr.h"
+#include "oneflow/core/framework/op_expr_helper.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+struct FusedScaleMaskSoftmaxDropoutInterState : public AutoGradCaptureState {
+  bool input_requires_grad = true;
+  float scale = 1.0;
+  float dropout_scale = 1.0;
+};
+
+class FusedScaleMaskSoftmaxDropout : public OpExprGradFunction<FusedScaleMaskSoftmaxDropoutInterState> {
+ public:
+  Maybe<void> Init(const OpExpr& op) override;
+  Maybe<void> Capture(FusedScaleMaskSoftmaxDropoutInterState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const FusedScaleMaskSoftmaxDropoutInterState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override;
+
+ private:
+  AttrMap base_attrs_;
+};
+
+Maybe<void> FusedScaleMaskSoftmaxDropout::Init(const OpExpr& op) {
+  const UserOpExpr* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
+  CHECK_NOTNULL_OR_RETURN(fw_op_expr);
+  base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> FusedScaleMaskSoftmaxDropout::Capture(FusedScaleMaskSoftmaxDropoutInterState* ctx,
+                                                  const TensorTuple& inputs, const TensorTuple& outputs,
+                                                  const AttrMap& attrs) const {
+  CHECK_EQ_OR_RETURN(inputs.size(), 3); // input, mask, dropout_mask
+  ctx->input_requires_grad = inputs.at(0)->requires_grad();
+
+  if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
+  ComposedAttrMap composed_attrs(attrs, base_attrs_);
+  ctx->scale = JUST(composed_attrs.GetAttr<float>("scale_value"));
+  ctx->dropout_scale = JUST(composed_attrs.GetAttr<float>("dropout_scale_value"));
+
+  ctx->SaveTensorForBackward(inputs.at(1));     // mask
+  ctx->SaveTensorForBackward(inputs.at(2));     // dropout_mask
+  ctx->SaveTensorForBackward(outputs.at(1));    // softmax_y
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> FusedScaleMaskSoftmaxDropout::Apply(const FusedScaleMaskSoftmaxDropoutInterState* ctx,
+                                                const TensorTuple& out_grads, TensorTuple* in_grads) const {
+  CHECK_EQ_OR_RETURN(out_grads.size(), 2);  // dy, d_softmax_y
+  if (!ctx->input_requires_grad) { return Maybe<void>::Ok(); }
+  in_grads->resize(3);  // input, mask, dropout_mask
+
+  const std::shared_ptr<oneflow::one::Tensor>& mask = ctx->SavedTensors().at(0);
+  const std::shared_ptr<oneflow::one::Tensor>& dropout_mask = ctx->SavedTensors().at(1);
+  const std::shared_ptr<oneflow::one::Tensor>& softmax_y = ctx->SavedTensors().at(2);
+  const std::shared_ptr<oneflow::one::Tensor>& input_grad = 
+      JUST(functional::FusedScaleMaskSoftmaxDropoutGrad(softmax_y, out_grads.at(0),
+          dropout_mask, mask, ctx->scale, ctx->dropout_scale));
+
+  in_grads->at(0) = input_grad;
+  return Maybe<void>::Ok();
+}
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("fused_scale_mask_softmax_dropout", FusedScaleMaskSoftmaxDropout);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/core/functional/functional_api.yaml

@@ -1487,6 +1487,22 @@
   signature: "Tensor (Tensor a, Tensor b, *, Float p=0.5, Int32 axis, Generator generator=None) => FusedBiasAddDropout"
   bind_python: True
 
+- name: "fused_scale_mask_softmax"
+  signature: "Tensor (Tensor x, Tensor mask, *, Float fill_value=0.0, Float scale=1.0) => FusedScaleMaskSoftmax"
+  bind_python: True
+
+- name: "fused_scale_mask_softmax_grad"
+  signature: "Tensor (Tensor y, Tensor dy, Tensor mask, Float scale=1.0) => FusedScaleMaskSoftmaxGrad"
+  bind_python: False
+
+- name: "fused_scale_mask_softmax_dropout"
+  signature: "TensorTuple (Tensor x, Tensor mask, *, Float fill_value=0.0, Float scale=1.0, Float p=0.5, Bool training=True, Generator generator=None) => FusedScaleMaskSoftmaxDropout"
+  bind_python: True
+
+- name: "fused_scale_mask_softmax_dropout_grad"
+  signature: "Tensor (Tensor softmax_y, Tensor dy, Tensor dropout_mask, Tensor mask, Float scale=1.0, Float dropout_scale=1.0) => FusedScaleMaskSoftmaxDropoutGrad"
+  bind_python: False
+
 - name: "fused_scale_tril_softmax_mask_scale"
   signature: "TensorTuple (Tensor a, *, Float p=0.5, Int64 diagonal, Float tril_scale_value, Generator generator=None) => FusedScaleTrilSoftmaxMaskScale"
   bind_python: True

oneflow/core/functional/impl/nn_functor.cpp

@@ -1886,6 +1886,68 @@ class FusedScaleTrilFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
+class FusedScaleMaskSoftmaxFunctor {
+ public:
+  FusedScaleMaskSoftmaxFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("fused_scale_mask_softmax")
+                        .Input("x")
+                        .Input("mask")
+                        .Output("y")
+                        .Build());
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const std::shared_ptr<one::Tensor>& mask,
+                           const float& fill_value, const float& scale) const {
+    MutableAttrMap attrs_;
+    JUST(attrs_.SetAttr<float>("scale_value", scale));
+    JUST(attrs_.SetAttr<float>("mask_fill_value", fill_value));
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {x, mask}, attrs_);
+  }
+ private:
+  std::shared_ptr<OpExpr> op_;
+};
+
+class FusedScaleMaskSoftmaxDropoutFunctor {
+ public:
+  FusedScaleMaskSoftmaxDropoutFunctor() {
+    random_mask_like_op_ = CHECK_JUST(one::OpBuilder("random_mask_like").Input("like").Output("out").Build());
+    fused_scale_mask_softmax_dropout_op_ = CHECK_JUST(one::OpBuilder("fused_scale_mask_softmax_dropout")
+                                                          .Input("x")
+                                                          .Input("mask")
+                                                          .Input("dropout_mask")
+                                                          .Output("y")
+                                                          .Output("softmax_y")
+                                                          .Build());
+  }
+  Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x, const std::shared_ptr<one::Tensor>& mask,
+                                const float& fill_value, const float& scale, const float& p, const bool& training, 
+                                const Optional<one::Generator>& generator) const {
+    float rate = p;
+    if (!training) rate = 0.0;
+    const auto gen = generator.value_or(JUST(one::DefaultAutoGenerator()));
+    MutableAttrMap random_mask_like_attrs;
+    JUST(random_mask_like_attrs.SetAttr<float>("rate", rate));
+    JUST(random_mask_like_attrs.SetAttr<int64_t>("seed", gen->current_seed()));
+    const auto& random_mask_like_state = std::make_shared<RandomMaskLikeKernelState>(gen);
+
+    const auto& dropout_mask = JUST(OpInterpUtil::Dispatch<Tensor>(
+        *random_mask_like_op_, {x},
+        OpExprInterpContext(random_mask_like_attrs, random_mask_like_state)));
+
+    float dropout_scale = 1.0;
+    if (rate != 1.0) { dropout_scale = 1.0 / (1.0 - rate); }
+    MutableAttrMap fused_scale_mask_softmax_dropout_attrs;
+    JUST(fused_scale_mask_softmax_dropout_attrs.SetAttr<float>("scale_value", scale));
+    JUST(fused_scale_mask_softmax_dropout_attrs.SetAttr<float>("mask_fill_value", fill_value));
+    JUST(fused_scale_mask_softmax_dropout_attrs.SetAttr<float>("dropout_scale_value", dropout_scale));
+
+    return OpInterpUtil::Dispatch<TensorTuple>(*fused_scale_mask_softmax_dropout_op_, 
+              {x, mask, dropout_mask}, fused_scale_mask_softmax_dropout_attrs);
+  }
+ private:
+  std::shared_ptr<OpExpr> random_mask_like_op_;
+  std::shared_ptr<OpExpr> fused_scale_mask_softmax_dropout_op_;
+};
+
 class CtcGreedyDecoderFunctor {
  public:
   CtcGreedyDecoderFunctor() {
@@ -2008,6 +2070,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::FusedBiasAddGeluFunctor>("FusedBiasAddGelu");
   m.add_functor<impl::FusedBiasAddGeluGradFunctor>("FusedBiasAddGeluGrad");
   m.add_functor<impl::FusedBiasAddDropoutFunctor>("FusedBiasAddDropout");
+  m.add_functor<impl::FusedScaleMaskSoftmaxFunctor>("FusedScaleMaskSoftmax");
+  m.add_functor<impl::FusedScaleMaskSoftmaxDropoutFunctor>("FusedScaleMaskSoftmaxDropout");
   m.add_functor<impl::FusedScaleTrilSoftmaxMaskScaleFunctor>("FusedScaleTrilSoftmaxMaskScale");
   m.add_functor<impl::FusedScaleTrilFunctor>("FusedScaleTril");
   m.add_functor<impl::CtcGreedyDecoderFunctor>("CtcGreedyDecoder");

oneflow/core/functional/impl/nn_grad_functor.cpp

@@ -806,6 +806,50 @@ class FusedScaleTrilSoftmaxMaskScaleGradFunctor {
   std::shared_ptr<OpExpr> fused_op_;
 };
 
+class FusedScaleMaskSoftmaxGradFunctor {
+ public:
+  FusedScaleMaskSoftmaxGradFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("fused_scale_mask_softmax_grad")
+                        .Input("y")
+                        .Input("dy")
+                        .Input("mask")
+                        .Output("dx")
+                        .Build());
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& y, const std::shared_ptr<one::Tensor>& dy,
+                           const std::shared_ptr<one::Tensor>& mask, const float& scale) const {
+    MutableAttrMap attrs_;
+    JUST(attrs_.SetAttr<float>("scale_value", scale));
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {y, dy, mask}, attrs_);
+  }
+ private:
+  std::shared_ptr<OpExpr> op_;
+};
+
+class FusedScaleMaskSoftmaxDropoutGradFunctor {
+ public:
+  FusedScaleMaskSoftmaxDropoutGradFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("fused_scale_mask_softmax_dropout_grad")
+                         .Input("softmax_y")
+                         .Input("dy")
+                         .Input("mask")
+                         .Input("dropout_mask")
+                         .Output("dx")
+                         .Build());
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& softmax_y, const std::shared_ptr<one::Tensor>& dy,
+                           const std::shared_ptr<one::Tensor>& dropout_mask, const std::shared_ptr<one::Tensor>& mask,
+                           const float& scale, const float& dropout_scale) const {
+    MutableAttrMap attrs_;
+    JUST(attrs_.SetAttr<float>("scale_value", scale));
+    JUST(attrs_.SetAttr<float>("dropout_scale_value", dropout_scale));
+
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {softmax_y, dy, dropout_mask, mask}, attrs_);
+  }
+ private:
+  std::shared_ptr<OpExpr> op_;
+};
+
 }  // namespace impl
 
 ONEFLOW_FUNCTION_LIBRARY(m) {
@@ -836,6 +880,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::BroadcastMatmulGradBFunctor>("BroadcastMatmulGradB");
   m.add_functor<impl::FusedScaleTrilSoftmaxMaskScaleGradFunctor>(
       "FusedScaleTrilSoftmaxMaskScaleGrad");
+  m.add_functor<impl::FusedScaleMaskSoftmaxGradFunctor>("FusedScaleMaskSoftmaxGrad");
+  m.add_functor<impl::FusedScaleMaskSoftmaxDropoutGradFunctor>("FusedScaleMaskSoftmaxDropoutGrad");
 };
 
 }  // namespace functional

oneflow/core/job_rewriter/auto_mixed_precision_lists.cpp

@@ -54,6 +54,8 @@ const AMPList& AutoMixedPrecisionLists::GrayList() {
                               "normalization_add_relu",
                               "sparse_softmax_cross_entropy",
                               "fused_tril_scale_softmax_mask_scale",
+                              "fused_scale_mask_softmax_dropout",
+                              "fused_scale_mask_softmax",
                               "fused_bias_add_gelu",
                               "fused_bias_add_mask_scale",
                               "acc"};