Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use op schema for cumsum #7175

Merged
merged 4 commits into from
Jan 4, 2022
Merged

use op schema for cumsum #7175

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0777968
add op schema for cumsum
wyushun Jan 4, 2022
3444a5a
change cumsum's td definition to math group
wyushun Jan 4, 2022
0932a3f
update
wyushun Jan 4, 2022
354b372
Merge branch 'master' into wyushun_cumsum_opschema
wyushun Jan 4, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion oneflow/core/functional/impl/math_functor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1687,7 +1687,7 @@ class MovedimIntFunctor {

class CumsumFunctor {
public:
CumsumFunctor() { op_ = CHECK_JUST(one::OpBuilder("cumsum").Input("in").Output("out").Build()); }
CumsumFunctor() { op_ = CHECK_JUST(one::OpBuilder("cumsum").Input("x").Output("y").Build()); }
Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& input, int64_t dim) const {
auto ndim = input->ndim();
if (dim < 0) { dim += ndim; }
Expand Down
32 changes: 32 additions & 0 deletions oneflow/ir/include/OneFlow/OneFlowUserOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3975,6 +3975,38 @@ def OneFlow_XlogyYGradOp : OneFlow_BaseOp<"xlogy_y_grad", [NoSideEffect, Declare
let has_data_type_infer_fn = 1;
}

def OneFlow_CumsumOp : OneFlow_BaseOp<"cumsum", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
let input = (ins
OneFlow_Tensor:$x
);
let output = (outs
OneFlow_Tensor:$y
);
let attrs = (ins
SI64Attr:$dim
);
let has_logical_tensor_desc_infer_fn = 1;
let has_physical_tensor_desc_infer_fn = 1;
let has_get_sbp_fn = 1;
let has_data_type_infer_fn = 1;
}

def OneFlow_CumsumGradOp : OneFlow_BaseOp<"cumsum_grad", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
let input = (ins
OneFlow_Tensor:$dy
);
let output = (outs
OneFlow_Tensor:$dx
);
let attrs = (ins
SI64Attr:$dim
);
let has_logical_tensor_desc_infer_fn = 1;
let has_physical_tensor_desc_infer_fn = 1;
let has_get_sbp_fn = 1;
let has_data_type_infer_fn = 1;
}

#endif // GET_ONEFLOW_MATH_OP_DEFINITIONS

// Group: MATMUL
Expand Down
6 changes: 3 additions & 3 deletions oneflow/user/kernels/cumsum_kernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,12 +61,12 @@ class CpuCumsumKernel final : public user_op::OpKernel {

private:
void Compute(user_op::KernelComputeContext* ctx) const override {
const auto* in = ctx->Tensor4ArgNameAndIndex("in", 0);
const auto* in = ctx->Tensor4ArgNameAndIndex("x", 0);
auto elem_cnt = in->shape().elem_cnt();
// judge whether tensor has 0 size dimension first
if (!elem_cnt) { return; }

auto* out = ctx->Tensor4ArgNameAndIndex("out", 0);
auto* out = ctx->Tensor4ArgNameAndIndex("y", 0);
auto dim = ctx->Attr<int64_t>("dim");
const auto* in_ptr = in->dptr<T>();
auto* out_ptr = out->mut_dptr<T>();
Expand All @@ -86,7 +86,7 @@ class CpuCumsumKernel final : public user_op::OpKernel {
#define REGISTER_CUMSUM_KERNEL(dtype) \
REGISTER_USER_KERNEL("cumsum").SetCreateFn<CpuCumsumKernel<dtype>>().SetIsMatchedHob( \
(user_op::HobDeviceType() == DeviceType::kCPU) \
&& (user_op::HobDataType("out", 0) == GetDataType<dtype>::value));
&& (user_op::HobDataType("y", 0) == GetDataType<dtype>::value));

REGISTER_CUMSUM_KERNEL(int64_t)
REGISTER_CUMSUM_KERNEL(float)
Expand Down
6 changes: 3 additions & 3 deletions oneflow/user/kernels/cumsum_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,11 +120,11 @@ class GpuCumsumKernel final : public user_op::OpKernel {
using user_op::OpKernel::Compute;
void Compute(user_op::KernelComputeContext* ctx) const override {
// judge whether tensor has 0 size dimension first
const auto* in = ctx->Tensor4ArgNameAndIndex("in", 0);
const auto* in = ctx->Tensor4ArgNameAndIndex("x", 0);
auto elem_cnt = in->shape().elem_cnt();
if (!elem_cnt) { return; }

auto* out = ctx->Tensor4ArgNameAndIndex("out", 0);
auto* out = ctx->Tensor4ArgNameAndIndex("y", 0);
auto dim = ctx->Attr<int64_t>("dim");
const auto* in_ptr = in->dptr<T>();
auto* out_ptr = out->mut_dptr<T>();
Expand Down Expand Up @@ -153,7 +153,7 @@ class GpuCumsumKernel final : public user_op::OpKernel {
#define REGISTER_CUDA_CUMSUM_KERNEL(dtype) \
REGISTER_USER_KERNEL("cumsum").SetCreateFn<GpuCumsumKernel<dtype>>().SetIsMatchedHob( \
(user_op::HobDeviceType() == DeviceType::kCUDA) \
&& (user_op::HobDataType("out", 0) == GetDataType<dtype>::value));
&& (user_op::HobDataType("y", 0) == GetDataType<dtype>::value));

REGISTER_CUDA_CUMSUM_KERNEL(int64_t)
REGISTER_CUDA_CUMSUM_KERNEL(float)
Expand Down
104 changes: 49 additions & 55 deletions oneflow/user/ops/cumsum_op.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,74 +14,68 @@ See the License for the specific language governing permissions and
limitations under the License.
*/
#include "oneflow/core/framework/framework.h"
#include "oneflow/core/framework/op_generated.h"

namespace oneflow {

namespace {
Maybe<void> CumsumOp::InferLogicalTensorDesc(user_op::InferContext* ctx) {
*ctx->OutputShape("y", 0) = ctx->InputShape("x", 0);
return Maybe<void>::Ok();
}

REGISTER_USER_OP("cumsum")
.Input("in")
.Output("out")
.Attr<int64_t>("dim")
.SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
*ctx->OutputShape("out", 0) = ctx->InputShape("in", 0);
return Maybe<void>::Ok();
})
.SetGetSbpFn([](user_op::SbpContext* ctx) -> Maybe<void> {
const auto& in_tensor_desc = ctx->LogicalTensorDesc4InputArgNameAndIndex("in", 0);
auto dim = ctx->Attr<int64_t>("dim");
for (auto i = dim + 1; i < in_tensor_desc.shape().NumAxes(); i++) {
ctx->NewBuilder()
.Split(user_op::OpArg("in", 0), i)
.Split(user_op::OpArg("out", 0), i)
.Build();
}
return Maybe<void>::Ok();
})
.SetDataTypeInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
*ctx->OutputDType("out", 0) = ctx->InputDType("in", 0);
return Maybe<void>::Ok();
});
Maybe<void> CumsumOp::InferPhysicalTensorDesc(user_op::InferContext* ctx) {
return InferLogicalTensorDesc(ctx);
}

REGISTER_USER_OP("cumsum_grad")
.Input("dy")
.Output("dx")
.Attr<int64_t>("dim")
.SetTensorDescInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
*ctx->OutputShape("dx", 0) = ctx->InputShape("dy", 0);
return Maybe<void>::Ok();
})
.SetGetSbpFn([](user_op::SbpContext* ctx) -> Maybe<void> {
const auto& dy_tensor_desc = ctx->LogicalTensorDesc4InputArgNameAndIndex("dy", 0);
for (auto i = 0; i < dy_tensor_desc.shape().NumAxes(); i++) {
ctx->NewBuilder()
.Split(user_op::OpArg("dy", 0), i)
.Split(user_op::OpArg("dx", 0), i)
.Build();
}
return Maybe<void>::Ok();
})
.SetDataTypeInferFn([](user_op::InferContext* ctx) -> Maybe<void> {
*ctx->OutputDType("dx", 0) = ctx->InputDType("dy", 0);
return Maybe<void>::Ok();
});
Maybe<void> CumsumOp::GetSbp(user_op::SbpContext* ctx) {
const auto& in_tensor_desc = ctx->LogicalTensorDesc4InputArgNameAndIndex("x", 0);
auto dim = ctx->Attr<int64_t>("dim");
for (auto i = dim + 1; i < in_tensor_desc.shape().NumAxes(); i++) {
ctx->NewBuilder().Split(user_op::OpArg("x", 0), i).Split(user_op::OpArg("y", 0), i).Build();
}
return Maybe<void>::Ok();
}

Maybe<void> CumsumOp::InferDataType(user_op::InferContext* ctx) {
*ctx->OutputDType("y", 0) = ctx->InputDType("x", 0);
return Maybe<void>::Ok();
}

Maybe<void> CumsumGradOp::InferLogicalTensorDesc(user_op::InferContext* ctx) {
*ctx->OutputShape("dx", 0) = ctx->InputShape("dy", 0);
return Maybe<void>::Ok();
}

Maybe<void> CumsumGradOp::InferPhysicalTensorDesc(user_op::InferContext* ctx) {
return InferLogicalTensorDesc(ctx);
}

Maybe<void> CumsumGradOp::GetSbp(user_op::SbpContext* ctx) {
const auto& dy_tensor_desc = ctx->LogicalTensorDesc4InputArgNameAndIndex("dy", 0);
for (auto i = 0; i < dy_tensor_desc.shape().NumAxes(); i++) {
ctx->NewBuilder().Split(user_op::OpArg("dy", 0), i).Split(user_op::OpArg("dx", 0), i).Build();
}
return Maybe<void>::Ok();
}

Maybe<void> CumsumGradOp::InferDataType(user_op::InferContext* ctx) {
*ctx->OutputDType("dx", 0) = ctx->InputDType("dy", 0);
return Maybe<void>::Ok();
}

REGISTER_USER_OP_GRAD("cumsum").SetGenBackwardOpConfFn(
[](const user_op::UserOpWrapper& op, const user_op::AddOpFn& AddOp) -> Maybe<void> {
if (op.NeedGenGradTensor4OpInput("in", 0)) {
if (op.NeedGenGradTensor4OpInput("x", 0)) {
user_op::UserOpConfWrapperBuilder builder(op.op_name() + "_grad");
user_op::UserOpConfWrapper grad_op =
builder.Op("cumsum_grad")
.Input("dy", op.GetGradTensorWithOpOutput("out", 0))
.Output("dx")
.Attr("dim", op.attr<int64_t>("dim"))
.Build();
user_op::UserOpConfWrapper grad_op = builder.Op("cumsum_grad")
.Input("dy", op.GetGradTensorWithOpOutput("y", 0))
.Output("dx")
.Attr("dim", op.attr<int64_t>("dim"))
.Build();
op.BindGradTensorWithOpInput(grad_op.output("dx", 0), "x", 0);
AddOp(grad_op);
}
return Maybe<void>::Ok();
});

} // namespace

} // namespace oneflow