| 24 | } |
| 25 | |
| 26 | std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible( |
| 27 | const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) { |
| 28 | auto&& op_def = def.cast_final_safe<Elemwise>(); |
| 29 | auto trait = megdnn::Elemwise::ModeTrait::from_mode(op_def.mode); |
| 30 | mgb_assert( |
| 31 | inputs.size() == trait.arity, "%s expects %u inputs; got %zu actually", |
| 32 | trait.name, trait.arity, inputs.size()); |
| 33 | TensorShapeArray inp_shapes; |
| 34 | DType out_dt; |
| 35 | CompNode out_cn; |
| 36 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 37 | auto&& t = inputs[i]; |
| 38 | if (!i) { |
| 39 | out_cn = t.comp_node; |
| 40 | out_dt = t.layout.dtype; |
| 41 | } else { |
| 42 | mgb_assert(t.comp_node == out_cn); |
| 43 | mgb_assert(t.layout.dtype == out_dt); |
| 44 | } |
| 45 | if (t.layout.ndim > 0) { |
| 46 | inp_shapes.push_back(t.layout); |
| 47 | } else { |
| 48 | TensorLayout out_layout; |
| 49 | out_layout.ndim = 0; |
| 50 | out_layout.dtype = out_dt; |
| 51 | return {{{out_layout, out_cn}}, false}; |
| 52 | } |
| 53 | } |
| 54 | // copy from megdnn::ElemwiseForward::check_dtype |
| 55 | switch (out_dt.category()) { |
| 56 | case DTypeCategory::FLOAT: |
| 57 | mgb_assert(trait.allow_float, "unsupport mode %s for float\n", trait.name); |
| 58 | break; |
| 59 | case DTypeCategory::INT: |
| 60 | mgb_assert(trait.allow_int, "unsupport mode %s for int\n", trait.name); |
| 61 | break; |
| 62 | case DTypeCategory::BOOL: |
| 63 | mgb_assert(trait.allow_bool, "unsupport mode %s for bool\n", trait.name); |
| 64 | break; |
| 65 | default: |
| 66 | // Quantized Dtype could also be handled by this op, |
| 67 | // but scales need to be the same. |
| 68 | break; |
| 69 | } |
| 70 | |
| 71 | auto&& out_shape = opr::Elemwise::get_output_var_shape(op_def.mode, inp_shapes); |
| 72 | return {{{TensorLayout(out_shape, out_dt, inputs[0].layout.format), out_cn}}, true}; |
| 73 | } |
| 74 | |
| 75 | DispatchMode decide_dispatch_mode( |
| 76 | const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) { |
nothing calls this directly
no test coverage detected