| 67 | } |
| 68 | |
| 69 | SmallVector<TensorPtr> apply_on_physical_tensor( |
| 70 | const OpDef& def, const SmallVector<TensorPtr>& inputs, |
| 71 | SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) { |
| 72 | // memory forward |
| 73 | if (memory_forward_success(def, inputs)) { |
| 74 | // maybe returns inputs[0] directly |
| 75 | return {Tensor::make( |
| 76 | inputs[0]->blob(), inputs[0]->offset(), inputs[0]->layout())}; |
| 77 | } |
| 78 | |
| 79 | if (inputs.size() == 2) { |
| 80 | // reduce to target shape, fallback to proxy_graph |
| 81 | return proxy_graph_detail::apply_on_physical_tensor( |
| 82 | def, inputs, output_descs, validated); |
| 83 | } |
| 84 | mgb_assert(inputs.size() == 1); |
| 85 | |
| 86 | auto comp_node = inputs[0]->comp_node(); |
| 87 | auto&& op_def = def.cast_final_safe<Reduce>(); |
| 88 | DnnOprCaller<megdnn::Reduce> dnn_op(comp_node, op_def.param()); |
| 89 | auto&& mode = dnn_op.param().mode; |
| 90 | auto& axis = dnn_op.param().axis; |
| 91 | auto keepdim = op_def.keepdim; |
| 92 | |
| 93 | DnnTensorND dnn_input = [&] { |
| 94 | if (axis == INT_MAX) { // reduce to scalar |
| 95 | axis = 0; |
| 96 | // flatten input |
| 97 | return inputs[0]->dnn_tensor({inputs[0]->shape().total_nr_elems()}); |
| 98 | } else { |
| 99 | if (axis < 0) { |
| 100 | axis = inputs[0]->layout().ndim + axis; |
| 101 | } |
| 102 | mgb_assert(axis >= 0 && axis < inputs[0]->layout().ndim); |
| 103 | return inputs[0]->dnn_tensor(); |
| 104 | } |
| 105 | }(); |
| 106 | auto output_layout = dnn_op.deduce_layout(dnn_input.layout); |
| 107 | auto resolve_keepdim = [&] { |
| 108 | if (!keepdim) { |
| 109 | if (output_layout.ndim > 1) { |
| 110 | mgb_assert(output_layout.shape[axis] == 1); |
| 111 | output_layout.remove_axis_inplace(axis); |
| 112 | } |
| 113 | } |
| 114 | }; |
| 115 | |
| 116 | TensorPtr output; |
| 117 | if (output_layout.is_empty()) { |
| 118 | // output empty, no computation |
| 119 | resolve_keepdim(); |
| 120 | output = Tensor::make(output_layout, comp_node); |
| 121 | } else if (dnn_input.layout.is_empty()) { |
| 122 | // input empty but output not, do fill |
| 123 | resolve_keepdim(); |
| 124 | output = Tensor::make(output_layout, comp_node); |
| 125 | auto on_bad_empty_reduce = [](const char* name, size_t axis) { |
| 126 | mgb_throw( |
nothing calls this directly
no test coverage detected