Browse Source

fix(mgb/opr): add layout constraint on the input of Cumsum

GitOrigin-RevId: d2d108e278
tags/v1.3.1
Megvii Engine Team 4 years ago
parent
commit
bab0caab72
2 changed files with 5 additions and 0 deletions
  1. +4
    -0
      src/opr/impl/misc.cpp
  2. +1
    -0
      src/opr/include/megbrain/opr/misc.h

+ 4
- 0
src/opr/impl/misc.cpp View File

@@ -132,6 +132,10 @@ void Cumsum::scn_do_execute() {
intl::get_megdnn_workspace_from_var(output().back()));
}

void Cumsum::add_input_layout_constraint() {
input(0)->add_layout_constraint_contiguous();
}

void Cumsum::init_output_static_infer_desc() {
using namespace cg::static_infer;
auto infer_shape = [](TensorShape& dest, const InpVal& iv) {


+ 1
- 0
src/opr/include/megbrain/opr/misc.h View File

@@ -86,6 +86,7 @@ public:
//! cumulative sum along given axis
MGB_DEFINE_OPR_CLASS(Cumsum, cg::SingleCNOperatorNodeBaseT<
mixin::MegDNNOprHolderImpl<megdnn::Cumsum>>) // {
void add_input_layout_constraint() override;

public:
Cumsum(VarNode *src, const Param &param, const OperatorNodeConfig &config);


Loading…
Cancel
Save