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

GitOrigin-RevId: d2d108e2786806aa3f4c44ebfc0bf5f094cdab42
jiangxiaobin96 · Apr 8, 2021 · bab0caa · bab0caa
1 parent 0d16539
commit bab0caa
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/src/opr/impl/misc.cpp b/src/opr/impl/misc.cpp
@@ -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) {

diff --git a/src/opr/include/megbrain/opr/misc.h b/src/opr/include/megbrain/opr/misc.h
@@ -86,6 +86,7 @@ MGB_DEFINE_OPR_CLASS(ArgsortBackward,
 //! 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);