| 25 | |
| 26 | namespace { |
| 27 | void check_opr_not_cross_mem(OperatorNodeBase* opr) { |
| 28 | if (opr->node_prop().contain( |
| 29 | OperatorNodeBase::NodeProp::Flag::CROSS_COMP_NODE_MEMORY)) |
| 30 | return; |
| 31 | MemNode mem_node_id; |
| 32 | bool first = true; |
| 33 | auto check = [&](VarNode* var) { |
| 34 | auto cur = var->comp_node().mem_node(); |
| 35 | mgb_assert(cur); |
| 36 | if (first) { |
| 37 | first = false; |
| 38 | mem_node_id = cur; |
| 39 | } else |
| 40 | mgb_assert( |
| 41 | mem_node_id == cur, |
| 42 | "for non cross-memory oprs, " |
| 43 | "all vars should reside on the same memory node"); |
| 44 | }; |
| 45 | for (auto i : opr->input()) { |
| 46 | check(i); |
| 47 | } |
| 48 | for (auto i : opr->output()) { |
| 49 | check(i); |
| 50 | } |
| 51 | } |
| 52 | |
| 53 | void update_output_shapes( |
| 54 | static_infer::StaticInferManagerImpl& infer_mgr, OperatorNodeBase* opr, |