Add Diagonal Op

docs/source/oneflow.rst

@@ -42,6 +42,7 @@ oneflow
             cos, 
             cosh, 
             diag, 
+            diagonal,
             movedim,
             div, 
             dot, 

docs/source/tensor.rst

@@ -42,6 +42,7 @@ OneFlow Tensor Class
             detach, 
             device, 
             diag, 
+            diagonal,
             dim, 
             div, 
             double, 

oneflow/core/autograd/gradient_funcs/diagonal.cpp

@@ -0,0 +1,72 @@
+/*
+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/attr_map.h"
+#include "oneflow/core/framework/op_expr_grad_function.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+struct DiagonalInterpState : public AutoGradCaptureState {
+  bool requires_grad = false;
+  int32_t offset = 0;
+};
+
+class Diagonal : public OpExprGradFunction<DiagonalInterpState> {
+ public:
+  Maybe<void> Init(const OpExpr& op) override;
+  Maybe<void> Capture(DiagonalInterpState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override;
+  Maybe<void> Apply(const DiagonalInterpState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override;
+
+ private:
+  AttrMap base_attrs_;
+};
+
+Maybe<void> Diagonal::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> Diagonal::Capture(DiagonalInterpState* ctx, const TensorTuple& inputs,
+                              const TensorTuple& outputs, const AttrMap& attrs) const {
+  CHECK_EQ_OR_RETURN(outputs.size(), 1);
+  ctx->requires_grad = inputs.at(0)->requires_grad();
+  if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
+  ComposedAttrMap composed_attrs(attrs, base_attrs_);
+  ctx->offset = JUST(composed_attrs.GetAttr<int32_t>("offset"));
+  ctx->SaveTensorForBackward(inputs.at(0));
+  return Maybe<void>::Ok();
+}
+
+Maybe<void> Diagonal::Apply(const DiagonalInterpState* ctx, const TensorTuple& out_grads,
+                            TensorTuple* in_grads) const {
+  CHECK_EQ_OR_RETURN(out_grads.size(), 1);
+  in_grads->resize(2);
+  if (ctx->requires_grad) {
+    const auto& x = ctx->SavedTensors().at(0);
+    in_grads->at(0) = JUST(functional::DiagonalGrad(out_grads.at(0), x, ctx->offset));
+  }
+  return Maybe<void>::Ok();
+}
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("diagonal", Diagonal);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/core/functional/functional_api.yaml

@@ -1456,6 +1456,14 @@
   signature: "Tensor (Tensor dy, Tensor in, Int32 diagonal=0) => DiagGrad"
   bind_python: False
 
+- name: "diagonal"
+  signature: "Tensor (Tensor x, Int32 offset=0, Int32 dim1=0, Int32 dim2=1) => Diagonal"
+  bind_python: True
+
+- name: "diagonal_grad"
+  signature: "Tensor (Tensor dy, Tensor in, Int32 offset=0) => DiagonalGrad"
+  bind_python: False
+
 - name: "tensor_getitem"
   signature: "Tensor (Tensor x, TensorIndex index) => TensorGetItem"
   bind_python: True

oneflow/core/functional/impl/array_functor.cpp

@@ -1590,6 +1590,63 @@ class DiagGradFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
+class DiagonalFunctor {
+ public:
+  DiagonalFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("diagonal").Input("in").Output("out").Build());
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const int32_t& offset,
+                           const int32_t& dim1, const int32_t& dim2) const {
+    int64_t ndims = x->shape()->NumAxes();
+
+    CHECK_GE_OR_RETURN(dim1, -ndims)
+        << ", Dimension out of range (expected to be in range of [" << -ndims << ", " << ndims - 1
+        << "], but got " << dim1 << ");";
+    CHECK_LT_OR_RETURN(dim1, ndims) << ", Dimension out of range (expected to be in range of ["
+                                    << -ndims << ", " << ndims - 1 << "], but got " << dim1 << ");";
+    CHECK_GE_OR_RETURN(dim2, -ndims)
+        << ", Dimension out of range (expected to be in range of [" << -ndims << ", " << ndims - 1
+        << "], but got " << dim2 << ");";
+    CHECK_LT_OR_RETURN(dim2, ndims) << ", Dimension out of range (expected to be in range of ["
+                                    << -ndims << ", " << ndims - 1 << "], but got " << dim2 << ");";
+
+    int32_t p_dim1 = dim1 >= 0 ? dim1 : dim1 + ndims;
+    int32_t p_dim2 = dim2 >= 0 ? dim2 : dim2 + ndims;
+    CHECK_NE_OR_RETURN(p_dim1, p_dim2)
+        << ", diagonal dimensions cannot be identical " << dim1 << ", " << dim2;
+
+    std::vector<int32_t> input_index{p_dim1, p_dim2};
+    for (int32_t i = 0; i < ndims; i++) {
+      if (i != p_dim1 && i != p_dim2) { input_index.push_back(i); }
+    }
+
+    std::shared_ptr<one::Tensor> d_x = JUST(Transpose(x, input_index));
+
+    MutableAttrMap attrs;
+    JUST(attrs.SetAttr<int32_t>("offset", offset));
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {d_x}, attrs);
+  }
+
+ private:
+  std::shared_ptr<OpExpr> op_;
+};
+
+class DiagonalGradFunctor {
+ public:
+  DiagonalGradFunctor() {
+    op_ = CHECK_JUST(one::OpBuilder("diagonal_grad").Input("dy").Input("in").Output("dx").Build());
+  }
+  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& dy,
+                           const std::shared_ptr<one::Tensor>& x, const int32_t& offset) const {
+    MutableAttrMap attrs;
+    JUST(attrs.SetAttr<int32_t>("offset", offset));
+    return OpInterpUtil::Dispatch<Tensor>(*op_, {dy, x}, attrs);
+  }
+
+ private:
+  std::shared_ptr<OpExpr> op_;
+};
+
 class TensorGetItemFunctor {
  public:
   TensorGetItemFunctor() {}
@@ -2306,6 +2363,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<impl::TriuFunctor>("Triu");
   m.add_functor<impl::DiagFunctor>("Diag");
   m.add_functor<impl::DiagGradFunctor>("DiagGrad");
+  m.add_functor<impl::DiagonalFunctor>("Diagonal");
+  m.add_functor<impl::DiagonalGradFunctor>("DiagonalGrad");
   m.add_functor<impl::TensorGetItemFunctor>("TensorGetItem");
   m.add_functor<impl::DimScatterFunctor>("DimScatter");
   m.add_functor<impl::DimScatterAddFunctor>("DimScatterAdd");

oneflow/user/kernels/diagonal_kernel.cpp

@@ -0,0 +1,139 @@
+/*
+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/common/util.h"
+#include "oneflow/core/framework/framework.h"
+#include "oneflow/core/kernel/new_kernel_util.h"
+#include "oneflow/core/kernel/kernel_util.h"
+#include "oneflow/core/ep/cuda/cuda_stream.h"
+
+
+namespace oneflow {
+namespace {
+
+template<typename T>
+struct DiagonalFunctor final {
+  void operator()(ep::Stream* stream, T* out_buf, const T* in_buf, int32_t size, int32_t dim1,
+                  int32_t dim2) {
+    int32_t offset_index = (dim1 + 1) * dim2;
+    FOR_RANGE(int32_t, index, 0, size * dim2) {
+      int32_t i = index / dim2;
+      int32_t j = index - i * dim2;
+      out_buf[j * size + i] = in_buf[i * offset_index + j];
+    }
+  }
+};
+
+template<typename T>
+struct DiagonalGradFunctor final {
+  void operator()(ep::Stream* stream, T* dx_buf, const T* dy_buf, int32_t size, int32_t dim1,
+                  int32_t dim2) {
+    int32_t offset_index = (dim1 + 1) * dim2;
+    FOR_RANGE(int32_t, index, 0, size * dim2) {
+      int32_t i = index / dim2;
+      int32_t j = index - i * dim2;
+      dx_buf[i * offset_index + j] = dy_buf[j * size + i];
+    }
+  }
+};
+
+}  // namespace
+
+template<typename T>
+class CpuDiagonalKernel final : public user_op::OpKernel {
+ public:
+  CpuDiagonalKernel() = default;
+  ~CpuDiagonalKernel() = default;
+
+ private:
+  void Compute(user_op::KernelComputeContext* ctx) const override {
+    const int32_t offset = ctx->Attr<int32_t>("offset");
+    const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0);
+    user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
+    const ShapeView& out_shape = out->shape();
+    const ShapeView& in_shape = in->shape();
+    const T* in_buf = in->dptr<T>();
+    T* out_buf = out->mut_dptr<T>();
+
+    int32_t size = out_shape.At(out_shape.NumAxes() - 1);
+    int32_t dim1 = in_shape.At(1);
+    int32_t dim2 = 0;
+    if (in_shape.NumAxes() <= 2) {
+      dim2 = 1;
+    } else {
+      dim2 = in_shape.Count(2, in_shape.NumAxes());
+    }
+
+    int32_t offset_in_bufer = (offset >= 0 ? offset * dim2 : -offset * dim1 * dim2);
+    in_buf += offset_in_bufer;
+    DiagonalFunctor<T>()(ctx->stream(), out_buf, in_buf, size, dim1, dim2);
+  }
+  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
+};
+
+template<typename T>
+class CpuDiagonalBackwardKernel final : public user_op::OpKernel {
+ public:
+  CpuDiagonalBackwardKernel() = default;
+  ~CpuDiagonalBackwardKernel() = default;
+
+ private:
+  void Compute(user_op::KernelComputeContext* ctx) const override {
+    const user_op::Tensor* dy = ctx->Tensor4ArgNameAndIndex("dy", 0);
+    user_op::Tensor* dx = ctx->Tensor4ArgNameAndIndex("dx", 0);
+    int32_t offset = ctx->Attr<int32_t>("offset");
+    const ShapeView& dx_shape = dx->shape();
+    const ShapeView& dy_shape = dy->shape();
+    T* dx_buf = dx->mut_dptr<T>();
+    const T* dy_buf = dy->dptr<T>();
+
+    Memset<DeviceType::kCPU>(ctx->stream(), dx->mut_dptr<T>(), 0, dx_shape.elem_cnt() * sizeof(T));
+
+    int32_t dim1 = dx_shape.At(1);
+    int32_t dim2 = 0;
+    if (dx_shape.NumAxes() <= 2) {
+      dim2 = 1;
+    } else {
+      dim2 = dx_shape.Count(2, dx_shape.NumAxes());
+    }
+    int32_t size = dy_shape.At(dy_shape.NumAxes() - 1);
+    int32_t offset_in_bufer = (offset >= 0 ? offset * dim2 : -offset * dim1 * dim2);
+    dx_buf += offset_in_bufer;
+
+    DiagonalGradFunctor<T>()(ctx->stream(), dx_buf, dy_buf, size, dim1, dim2);
+  }
+  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
+};
+
+#define REGISTER_DIAGONAL_KERNELS(dtype)                                        \
+  REGISTER_USER_KERNEL("diagonal")                                                      \
+      .SetCreateFn<CpuDiagonalKernel<dtype>>()                                     \
+      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCPU)                              \
+                       && (user_op::HobDataType("in", 0) == GetDataType<dtype>::value)); \
+  REGISTER_USER_KERNEL("diagonal_grad")                                                 \
+      .SetCreateFn<CpuDiagonalBackwardKernel<dtype>>()                             \
+      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCPU)                              \
+                       && (user_op::HobDataType("in", 0) == GetDataType<dtype>::value));
+
+REGISTER_DIAGONAL_KERNELS(float);
+REGISTER_DIAGONAL_KERNELS(double);
+REGISTER_DIAGONAL_KERNELS(int8_t);
+REGISTER_DIAGONAL_KERNELS(int32_t);
+REGISTER_DIAGONAL_KERNELS(int64_t);
+
+#undef REGISTER_DIAGONAL_KERNELS
+
+}  // namespace oneflow