Merge slice and logical slice

docs/source/oneflow.rst

@@ -136,7 +136,7 @@ oneflow
             selu, 
             silu, 
             slice, 
-            logical_slice,  
+            slice_update,
             softsign, 
             sort, 
             softplus, 

oneflow/core/autograd/gradient_funcs/slice.cpp

@@ -23,42 +23,50 @@ namespace oneflow {
 namespace one {
 
 struct SliceCaptureState : public AutoGradCaptureState {
-  bool requires_grad;
   Shape like_shape;
   std::vector<int64_t> start;
   std::vector<int64_t> stop;
   std::vector<int64_t> step;
+  Symbol<NdSbp> in_sbp;
 };
 
 class Slice : public OpExprGradFunction<SliceCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
-    CHECK_NOTNULL_OR_RETURN(fw_op_expr);
+    CHECK_NOTNULL_OR_RETURN(fw_op_expr) << "Slice op_expr is null";
     base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
     return Maybe<void>::Ok();
   }
 
   Maybe<void> Capture(SliceCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
                       const AttrMap& attrs) const override {
-    CHECK_EQ_OR_RETURN(inputs.size(), 1);
-    CHECK_EQ_OR_RETURN(outputs.size(), 1);
-    ctx->requires_grad = inputs.at(0)->requires_grad();
-    if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
+    CHECK_EQ_OR_RETURN(inputs.size(), 1) << "Slice input size must be 1";
+    CHECK_EQ_OR_RETURN(outputs.size(), 1) << "Slice output size must be 1";
 
     ComposedAttrMap composed_attrs(attrs, base_attrs_);
     ctx->start = JUST(composed_attrs.GetAttr<std::vector<int64_t>>("start"));
     ctx->stop = JUST(composed_attrs.GetAttr<std::vector<int64_t>>("stop"));
     ctx->step = JUST(composed_attrs.GetAttr<std::vector<int64_t>>("step"));
-    ctx->like_shape = *(inputs.at(0)->shape());
+    ctx->like_shape = *(inputs[0]->shape());
+    if (inputs[0]->is_consistent()) { ctx->in_sbp = JUST(inputs[0]->nd_sbp()); }
     return Maybe<void>::Ok();
   }
 
   Maybe<void> Apply(const SliceCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(1);
-    in_grads->at(0) = JUST(
-        functional::SliceGrad(out_grads.at(0), ctx->like_shape, ctx->start, ctx->stop, ctx->step));
+    std::shared_ptr<Tensor> zeros;
+    if (out_grads[0]->is_local()) {
+      zeros = JUST(functional::Constant(ctx->like_shape, 0, out_grads[0]->dtype(),
+                                        JUST(out_grads[0]->device())));
+    } else {
+      const auto& parallel_desc = JUST(out_grads[0]->parallel_desc());
+      zeros = JUST(functional::ConsistentConstant(ctx->like_shape, 0, out_grads[0]->dtype(),
+                                                  parallel_desc, *JUST(GetSbpList(ctx->in_sbp))));
+    }
+    (*in_grads)[0] = JUST(functional::SliceUpdate(zeros, out_grads[0], ctx->start, ctx->stop,
+                                                  ctx->step, /*inplace=*/false));
     return Maybe<void>::Ok();
   }
 
@@ -67,51 +75,65 @@ class Slice : public OpExprGradFunction<SliceCaptureState> {
 };
 
 struct SliceUpdateCaptureState : public AutoGradCaptureState {
-  bool requires_grad_x;
-  bool requires_grad_update;
+  bool requires_grad_ref = false;
+  bool requires_grad_value = false;
   std::vector<int64_t> start;
   std::vector<int64_t> stop;
   std::vector<int64_t> step;
+  Shape value_shape;  // used to calculate ref gradient
+  Symbol<NdSbp> value_sbp;
 };
 
 class SliceUpdate : public OpExprGradFunction<SliceUpdateCaptureState> {
  public:
   Maybe<void> Init(const OpExpr& op) override {
     const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
-    CHECK_NOTNULL_OR_RETURN(fw_op_expr);
+    CHECK_NOTNULL_OR_RETURN(fw_op_expr) << "SliceUpdate op_expr is null";
 
     base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
     return Maybe<void>::Ok();
   }
 
   Maybe<void> Capture(SliceUpdateCaptureState* ctx, const TensorTuple& inputs,
                       const TensorTuple& outputs, const AttrMap& attrs) const override {
-    CHECK_EQ_OR_RETURN(inputs.size(), 2);
-    CHECK_EQ_OR_RETURN(outputs.size(), 1);
-    ctx->requires_grad_x = inputs.at(0)->requires_grad();
-    ctx->requires_grad_update = inputs.at(1)->requires_grad();
-    if (!ctx->requires_grad_x && !ctx->requires_grad_update) { return Maybe<void>::Ok(); }
+    CHECK_EQ_OR_RETURN(inputs.size(), 2) << "SliceUpdate input size must be 2";
+    CHECK_EQ_OR_RETURN(outputs.size(), 1) << "SliceUpdate output size must be 1";
+    ctx->requires_grad_ref = inputs[0]->requires_grad();
+    ctx->requires_grad_value = inputs[1]->requires_grad();
+    if (!ctx->requires_grad_ref && !ctx->requires_grad_value) { return Maybe<void>::Ok(); }
 
     ComposedAttrMap composed_attrs(attrs, base_attrs_);
     ctx->start = JUST(composed_attrs.GetAttr<std::vector<int64_t>>("start"));
     ctx->stop = JUST(composed_attrs.GetAttr<std::vector<int64_t>>("stop"));
     ctx->step = JUST(composed_attrs.GetAttr<std::vector<int64_t>>("step"));
 
-    if (ctx->requires_grad_x) { ctx->SaveTensorForBackward(inputs.at(1)); }
+    if (ctx->requires_grad_ref) {
+      ctx->value_shape = *(inputs[1]->shape());
+      if (inputs[1]->is_consistent()) { ctx->value_sbp = JUST(inputs[1]->nd_sbp()); }
+    }
     return Maybe<void>::Ok();
   }
 
   Maybe<void> Apply(const SliceUpdateCaptureState* ctx, const TensorTuple& out_grads,
                     TensorTuple* in_grads) const override {
     in_grads->resize(2);
 
-    if (ctx->requires_grad_x) {
-      const auto& update = ctx->SavedTensors().at(0);
-      const auto& temp = JUST(functional::ZerosLike(update));
-      (*in_grads)[0] = JUST(functional::SliceUpdate(out_grads[0], temp, ctx->start, ctx->stop,
-                                                    ctx->step, /*inplace=*/false));
+    if (ctx->requires_grad_ref) {
+      std::shared_ptr<Tensor> zeros;
+      if (out_grads[0]->is_local()) {
+        zeros = JUST(functional::Constant(ctx->value_shape, 0, out_grads[0]->dtype(),
+                                          JUST(out_grads[0]->device())));
+      } else {
+        const auto& parallel_desc = JUST(out_grads[0]->parallel_desc());
+        zeros =
+            JUST(functional::ConsistentConstant(ctx->value_shape, 0, out_grads[0]->dtype(),
+                                                parallel_desc, *JUST(GetSbpList(ctx->value_sbp))));
+      }
+      (*in_grads)[0] =
+          JUST(functional::SliceUpdate(JUST(functional::Identity(out_grads[0])), zeros, ctx->start,
+                                       ctx->stop, ctx->step, /*inplace=*/false));
     }
-    if (ctx->requires_grad_update) {
+    if (ctx->requires_grad_value) {
       (*in_grads)[1] = JUST(functional::Slice(out_grads[0], ctx->start, ctx->stop, ctx->step,
                                               /*enable_view_slice=*/false));
     }
@@ -122,8 +144,8 @@ class SliceUpdate : public OpExprGradFunction<SliceUpdateCaptureState> {
   AttrMap base_attrs_;
 };
 
-REGISTER_OP_EXPR_GRAD_FUNCTION("slice", Slice);
 REGISTER_OP_EXPR_GRAD_FUNCTION("slice_update", SliceUpdate);
+REGISTER_OP_EXPR_GRAD_FUNCTION("slice", Slice);
 
 }  // namespace one
 }  // namespace oneflow

oneflow/core/boxing/symmetric_b_to_s_boxing.cpp

@@ -88,8 +88,8 @@ Maybe<one::Tensor> SymmetricB2S(const std::shared_ptr<one::Tensor>& tensor, Symb
       start.emplace_back(range.begin());
       stop.emplace_back(range.end());
     }
-    local_tensor =
-        JUST(one::functional::Slice(local_tensor, start, stop, step, /*enable_view_slice=*/false));
+    local_tensor = JUST(one::functional::Slice(local_tensor, start, stop, step,
+                                               /*enable_view_slice=*/false));
   }
 
   return JUST(one::functional::LocalToConsistent(local_tensor, out->placement(),

oneflow/core/framework/tensor_methods.cpp

@@ -152,15 +152,26 @@ Maybe<Tensor> Slice(const std::shared_ptr<Tensor>& input, const std::vector<int6
   }
 
   auto output = JUST(BasicView(input, Shape(target_dims), Stride(target_strides), storage_offset));
+  Symbol<NdSbp> in_nd_sbp;
+  if (input->is_consistent()) { in_nd_sbp = JUST(input->nd_sbp()); }
+
   if (autograd::GradMode::is_enabled() && input->requires_grad()) {
     auto backward_fn = std::make_shared<BackwardFunction>();
     backward_fn->body = [=](const TensorTuple& out_grads, TensorTuple* in_grads,
                             bool create_graph) -> Maybe<void> {
       autograd::AutoGradMode mode(create_graph);
       CHECK_EQ_OR_RETURN(out_grads.size(), 1);  // NOLINT(maybe-need-error-msg)
-      in_grads->resize(1);
-      (*in_grads)[0] = JUST(functional::SliceGrad(
-          JUST(VectorAt(out_grads, 0)), Shape(input->shape()->dim_vec()), starts, ends, steps));
+      std::shared_ptr<Tensor> zeros;
+      if (out_grads[0]->is_local()) {
+        zeros = JUST(
+            functional::Constant(*shape, 0, out_grads[0]->dtype(), JUST(out_grads[0]->device())));
+      } else {
+        const auto& parallel_desc = JUST(out_grads[0]->parallel_desc());
+        zeros = JUST(functional::ConsistentConstant(*shape, 0, out_grads[0]->dtype(), parallel_desc,
+                                                    *JUST(GetSbpList(in_nd_sbp))));
+      }
+      (*in_grads)[0] = JUST(
+          functional::SliceUpdate(zeros, out_grads[0], starts, ends, steps, /*inplace=*/false));
       return Maybe<void>::Ok();
     };
     backward_fn->status = []() { return true; };

oneflow/core/functional/functional_api.yaml

@@ -1297,14 +1297,6 @@
   signature: "Tensor (Tensor x, Int64 start, Int64 end) => SliceView1dContiguous"
   bind_python: True
 
-- name: "slice"
-  signature: "Tensor (Tensor x, Int64List start, Int64List stop, Int64List step, Bool enable_view_slice=None) => Slice"
-  bind_python: True
-
-- name: "slice_grad"
-  signature: "Tensor (Tensor dy, Shape like, Int64List start, Int64List stop, Int64List step) => SliceGrad"
-  bind_python: False
-
 - name: "narrow"
   signature: "Tensor (Tensor input, Int64 dim, Int64 start, Int64 length) => Narrow"
   bind_python: True
@@ -1313,16 +1305,12 @@
   signature: "Tensor (Tensor dy, Tensor like, Int64 dim, Int64 start, Int64 length) => NarrowGrad"
   bind_python: False
 
-- name: "slice_update"
-  signature: "Tensor (Tensor x, Tensor update, Int64List start, Int64List stop, Int64List step, *, Bool inplace=False) => SliceUpdate"
-  bind_python: True
-
-- name: "logical_slice"
-  signature: "Tensor (Tensor x, Int64List start, Int64List stop, Int64List step, Bool enable_view_slice=None) => LogicalSlice"
+- name: "slice"
+  signature: "Tensor (Tensor x, Int64List start, Int64List stop, Int64List step, Bool enable_view_slice=None) => Slice"
   bind_python: True
 
-- name: "logical_slice_assign"
-  signature: "Tensor (Tensor ref, Tensor value, Int64List start, Int64List stop, Int64List step) => LogicalSliceAssign"
+- name: "slice_update"
+  signature: "Tensor (Tensor ref, Tensor value, Int64List start, Int64List stop, Int64List step, Bool inplace=False) => SliceUpdate"
   bind_python: True
 
 - name: "copy"