| 11 | |
| 12 | template <typename Opr> |
| 13 | CompNode get_device(const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) { |
| 14 | auto&& op_def = def.cast_final_safe<Opr>(); |
| 15 | const char* op_name = op_def.make_name().c_str(); |
| 16 | CompNode oup_cn = op_def.comp_node; |
| 17 | if (!oup_cn.valid()) { |
| 18 | size_t nr_inp = inputs.size(); |
| 19 | mgb_assert( |
| 20 | nr_inp > 0, "number of inputs of %s should be greater than 0", op_name); |
| 21 | auto&& inp_cn = inputs[0].comp_node; |
| 22 | for (size_t i = 1; i < nr_inp; ++i) { |
| 23 | mgb_assert( |
| 24 | inp_cn == inputs[i].comp_node, |
| 25 | "input tensors of %s operator should have same device, but get " |
| 26 | "%s vs %s", |
| 27 | op_name, inp_cn.to_string().c_str(), |
| 28 | inputs[i].comp_node.to_string().c_str()); |
| 29 | } |
| 30 | oup_cn = inp_cn; |
| 31 | } |
| 32 | return oup_cn; |
| 33 | } |
| 34 | |
| 35 | bool is_all_inputs_valid(const SmallVector<LogicalTensorDesc>& inputs) { |
| 36 | bool input_valid = true; |